ShardingInfo

ShardingInfo#

class treescope.ndarray_adapters.ShardingInfo[source]#

Summary of the sharding of an array.

Variables:

shard_shape (tuple[int, ...]) – Shape of a single shard.
device_index_to_shard_slices (dict[int, tuple[slice, ...]]) – A mapping from device index to the tuple of per-axis slices of the original array that is assigned to that device. The length of each axis slice must match the shard_shape along that axis (or be the full slice slice(None)).
device_type (str) – The type of device that the array is sharded across, as a string (e.g. “CPU”, “TPU”, “GPU”).
fully_replicated (bool) – Whether the array is fully replicated across all devices.

Methods

__init__(shard_shape, ...[, fully_replicated])

Attributes