Python function

shard_shape

shard_shape()

max.experimental.sharding.shard_shape(global_shape, placements, mesh_shape)

source

Computes the per-shard shape from a global shape and placements.

For each Sharded(axis=k) placement on mesh axis i, dimension k is divided by mesh_shape[i]. Replicated and Partial placements leave the shape unchanged. Works with both static and symbolic dimensions.

Parameters:

• global_shape (Sequence[int | str | Dim | integer[Any] | TypedAttr]) – The global shape of the tensor.
• placements (Sequence[Placement]) – One Placement per mesh axis.
• mesh_shape (Sequence[int]) – The size of each mesh axis.

Returns:

The per-shard dimensions after applying every sharded placement.

Return type:

list[Dim]

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit