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.custom	User-facing definitions of custom ops for eager and graph use.
max.experimental.executor	Executor protocol and concrete implementations for eager graph execution.
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.random	Provides random tensor generation utilities.
max.experimental.realization_context	Implementations of various realization strategies.
max.experimental.sharding	Defines how a tensor is laid out across a device mesh and how ops reshard it.
max.experimental.sharding.per_shard_dim	Per-shard Dim wrapper carrying one value per device.
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.
max.experimental.validation	Guard rails for running eager code in production.

Custom operations

CustomOp	A declared custom op; call it eagerly or in a graph.