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. Should be the same length as the number of slices in each value of device_index_to_shard_slices.
device_index_to_shard_slices (dict[int, tuple[slice | int, ...]]) – A mapping from device index to the tuple of per-axis indices or slices of the original array that is assigned to that device. Each entry of this tuple should either be an int or a slice object. If an int, that axis should not appear in shard_shape (e.g. the full array is formed by stacking the shards along a new axis). If a slice, the corresponding axis should appear in shard_shape, and the slice should be the full slice slice(None) (if the array is not sharded over this axis) or a slice that matches the corresponding entry in shard_shape (if the full array is formed by concatenating the shards along this axis).
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