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

DeviceMapping

DeviceMapping

class max.tensor.sharding.DeviceMapping(mesh, placements)

source

Bases: object

How a tensor is distributed across a device mesh.

Parameters:

• mesh (DeviceMesh) – The device mesh.
• placements (tuple[Placement, ...]) – One Placement per mesh axis, in mesh-axis order.

check_shape()

check_shape(shape)

source

Raises if this mapping shards an axis shape does not have.

A negative axis counts from the end of shape, so Sharded(-1) names the last axis, matching how the rest of the sharding code normalizes axes. This checks range only: an axis that would split into empty shards passes here, and transfer_to() refuses it later.

Parameters:

shape (Sized) – The global shape to check against.

Raises:

ValueError – If a Sharded axis is below -len(shape) or at least len(shape).

Return type:

None

is_fully_replicated

property is_fully_replicated: bool

source

True when every mesh axis is Replicated.

mesh

mesh: DeviceMesh

source

placements

placements: tuple[Placement, ...]

source

replicated()

classmethod replicated(mesh)

source

Builds a mapping that replicates a tensor on every axis of mesh.

Parameters:

mesh (DeviceMesh) – The device mesh to replicate over.

Returns:

A mapping with Replicated on each of mesh’s axes.

Return type:

DeviceMapping