For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

TensorLayout

TensorLayout

class max.experimental.sharding.TensorLayout(dtype, shape, device)

source

Bases: object

A tensor’s dtype and global shape, distributed across a device mesh.

A compiled callable reads an argument given this layout but may not store through it. See BufferLayout for one it may write.

A SymbolicDim sharded along a mesh axis becomes one fresh local dim per shard, named "{original}_{axis_name}_{shard}", so sharding the same global dim on different axes stays distinguishable.

Parameters:

• dtype (DType) – The element data type.
• shape (Shape) – The global shape, not one shard’s. A negative Sharded axis counts from the end of it.
• device (Device | DeviceRef | DeviceMesh | DeviceMapping) – The placement: a single device or DeviceRef, a mesh to replicate over, or a DeviceMapping. It is stored as mapping; device is the single device a non-distributed value sits on.

Raises:

• TypeError – If device is none of the accepted types.
• ValueError – If the placement count differs from the mesh’s axis count, or a Sharded axis falls outside the shape’s dimensions.

as_buffer()

as_buffer()

source

Returns this layout as a buffer a callable may store through.

Return type:

BufferLayout

as_tensor()

as_tensor()

source

Returns this layout as a value a callable only reads.

Return type:

TensorLayout

device

property device: Device

source

The single device this value sits on.

Raises:

ValueError – If it spans more than one device.

dtype

dtype: DType

source

The element data type.

local_types

property local_types: Sequence[_TensorTypeBase[Any]]

source

One TensorType per device, in mesh order.

mapping

mapping: DeviceMapping

source

The distribution across the device mesh.

mesh

property mesh: DeviceMesh

source

The mesh this value is distributed over.

placements

property placements: tuple[Placement, ...]

source

One placement per mesh axis.

rank

property rank: int

source

The number of dimensions.

shape

shape: Shape

source

The global shape.