For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Python class
DeviceMesh
DeviceMesh
class max.experimental.sharding.DeviceMesh(devices, mesh_shape, axis_names)
Bases: object
An N-dimensional logical grid of devices.
axis_names
The human-readable name for each mesh axis.
axis_size()
axis_size(axis)
Returns the size of a mesh axis by name or index.
-
Raises:
-
- ValueError – If
axisis a name not on the mesh. - IndexError – If
axisis an integer outside[0, ndim).
- ValueError – If
-
Parameters:
-
Return type:
default()
static default()
Returns a single-device mesh on the default device (CPU).
-
Return type:
device_coord()
device_coord(device_idx, axis)
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.
-
Parameters:
-
Returns:
-
The device’s coordinate along axis, in
[0, axis_size). -
Raises:
-
- IndexError – If
device_idxis out of range, or ifaxisis an integer outside[0, ndim). - ValueError – If
axisis a name not present on the mesh.
- IndexError – If
-
Return type:
devices
The devices in the mesh, in row-major order.
is_simulated
property is_simulated: bool
Returns True if all mesh slots reference the same device.
A simulated mesh uses graph-level ops to emulate multi-device collectives on a single CPU or GPU.
is_single
property is_single: bool
Returns True if this mesh contains exactly one device.
mesh_shape
The shape of the logical grid.
ndim
property ndim: int
The number of mesh dimensions.
num_devices
property num_devices: int
The total number of devices.
single()
static single(device)
Creates a trivial single-device mesh.
-
Parameters:
-
device (Device)
-
Return type:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!