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)
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.
- device (Device | DeviceRef | DeviceMesh | DeviceMapping) – The placement: a single device, a mesh to replicate over,
or a
DeviceMapping.
-
Raises:
-
ValueError – If a placement shards an axis this shape does not have.
as_buffer()
as_buffer()
Returns this layout as a buffer a callable may store through.
-
Return type:
as_tensor()
as_tensor()
Returns this layout as a value a callable only reads.
-
Return type:
device
property device: Device
The single device this value sits on.
-
Raises:
-
ValueError – If it spans more than one device.
dtype
dtype: DType
The element data type.
local_types
One TensorType per device, in mesh order.
mapping
mapping: DeviceMapping
The distribution across the device mesh.
mesh
property mesh: DeviceMesh
The mesh this value is distributed over.
placements
One placement per mesh axis.
rank
property rank: int
The number of dimensions.
shape
shape: Shape
The global shape.