tensorplay.distributed.device_mesh.init_device_mesh

tensorplay.distributed.device_mesh.init_device_mesh(device_type: str, mesh_shape: tuple[int, ...], *, mesh_dim_names: tuple[str, ...] | None = None, backend_override=None) → DeviceMesh[source]

Initializes a DeviceMesh based on device_type and mesh_shape (torch parity).

This creates a DeviceMesh with an n-dimensional array layout, where n is the length of mesh_shape. If mesh_dim_names is provided, each dimension is labeled as mesh_dim_names[i].

Note

Follows SPMD: ensure mesh_shape is identical across all ranks.

Example:

>>> mesh_1d = init_device_mesh("cuda", mesh_shape=(8,))
>>> mesh_2d = init_device_mesh("cuda", mesh_shape=(2, 8),
...                             mesh_dim_names=("dp", "tp"))