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 function

mesh_context

`mesh_context()`

max.experimental.sharding.mesh_context(mesh)

source

Publishes mesh to spec-first NamedMapping constructions.

JAX-style: when a NamedMapping is created without an explicit mesh inside this block, it picks up mesh from this context and resolves the spec against it.

from max.driver import CPU
from max.experimental.sharding import (
    DeviceMesh,
    get_active_mesh,
    mesh_context,
)

mesh = DeviceMesh(
    devices=(CPU(), CPU()), mesh_shape=(2,), axis_names=("tp",)
)

with mesh_context(mesh):
    # Spec-first constructions inside this block resolve against
    # ``mesh`` without naming it explicitly.
    active = get_active_mesh()

Parameters:: mesh (DeviceMesh)
Return type:: Iterator[DeviceMesh]

mesh_context()​

`mesh_context()`