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.experimental.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.

Parameters:

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

Raises:

ValueError – If a Sharded axis is out of range.

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