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

DeviceMesh

DeviceMesh

class max.experimental.sharding.DeviceMesh(devices, mesh_shape, axis_names)

source

Bases: object

An N-dimensional logical grid of devices.

devices holds prod(mesh_shape) entries in row-major order. A mesh of real devices holds distinct entries; a simulated mesh repeats one device in every slot, and the graph computes its collectives instead of the hardware. axis_names holds one unique name per axis.

Raises:

ValueError – If devices is empty, its length differs from prod(mesh_shape), axis_names differs in length from mesh_shape or repeats a name, or devices mixes repeated and distinct entries.

Parameters:

• devices (tuple[Device, ...])
• mesh_shape (tuple[int, ...])
• axis_names (tuple[str, ...])

axis_names

axis_names: tuple[str, ...]

source

The human-readable name for each mesh axis.

axis_size()

axis_size(axis)

source

Returns the size of a mesh axis by name or index.

For several axes, returns the number of devices they span together, which is the product of their sizes. On a (2, 3) mesh, mesh.axis_size((0, 1)) == 6. Pass each axis at most once: a repeated axis contributes its size once per occurrence.

Parameters:

axis (str | int | Sequence[str | int]) – The mesh axis or axes, by name or integer index.

Returns:

The number of devices along axis.

Raises:

• ValueError – If axis is a name not on the mesh.
• IndexError – If axis is an integer outside [0, ndim).

Return type:

int

device_coord()

device_coord(device_idx, axis)

source

Returns device_idx’s coordinate along the given mesh axis.

For a mesh shaped (2, 3) with row-major device ordering, the device at flat index 4 has coords (1, 1) — so mesh.device_coord(4, 0) == 1 and mesh.device_coord(4, 1) == 1.

For several axes, returns the device’s index among the devices they span, counted in row-major order over the axes as given, so mesh.device_coord(4, (0, 1)) == 4.

Parameters:

• device_idx (int) – The flat device index in row-major order.
• axis (str | int | Sequence[str | int]) – The mesh axis or axes to query, by name or integer index.

Returns:

The device’s coordinate along axis, in [0, axis_size(axis)).

Raises:

• IndexError – If device_idx is out of range, or if axis is an integer outside [0, ndim).
• ValueError – If axis is a name not present on the mesh.

Return type:

int

devices

devices: tuple[Device, ...]

source

The devices in the mesh, in row-major order.

mesh_shape

mesh_shape: tuple[int, ...]

source

The shape of the logical grid.

ndim

property ndim: int

source

The number of mesh dimensions.

num_devices

property num_devices: int

source

The total number of devices.

single()

static single(device)

source

Creates a one-device mesh.

Its shape is (1,) and its single axis is named "_".

Parameters:

device (Device) – The device the mesh holds.

Return type:

DeviceMesh