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 package

max.experimental

Experimental APIs for building, sharding, and running ML workloads.

Built on top of max.graph and max.driver, in three layers that each consume the one below:

• nn – the Module base class plus ahead-of-time compilation to a CompiledCallable.
• functional – a one-function-per-op distributed dispatcher (F.matmul, F.add, …).
• sharding – placements, the device mesh, the action data model, a cost model, and the auto_reshard policy.

The distributed Tensor ties them together. The following example calls Tensor.ones() to build a 4x8 tensor and passes it to F.matmul along with its transpose:

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor.ones((4, 8))
y = F.matmul(x, x.T)

Caution

Experimental. The public API is not stable; names and module structure may change.

Modules

max.experimental.compilation	Traces and compiles Python functions over tensors.
max.experimental.functional	Distributed functional ops with explicit per-op SPMD dispatch.
max.experimental.nn	Module framework for max.experimental.
max.experimental.nn.norm	Normalization layers for MAX neural networks.
max.experimental.nn.rope	Positional embedding modules and functions.
max.experimental.sharding	Defines how a tensor is laid out across a device mesh and how ops reshard it.
max.experimental.tensor	Provides tensor operations with eager execution capabilities.
max.experimental.testing	Testing utilities for the experimental Tensor API.
max.experimental.torch	Bridge between PyTorch and MAX graphs.