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 module

max.experimental.compilation

APIs to trace and compile callables.

compile() turns a function over tensors into a compiled one in two calls: first pass a spec (dtype, shape, device) for each tensor argument, then call the result on real tensors. Arguments that are not tensors are fixed while tracing, so pass them the same way both times. stage() stops after tracing, for inspecting the graph.

from max.driver import CPU
from max.dtype import DType
from max.experimental import compilation
from max.experimental.sharding import TensorLayout
from max.experimental.tensor import Tensor

def step(x: Tensor, *, gain: float) -> Tensor:
    return x * gain

x_spec = TensorLayout(DType.float32, ["batch", 2], CPU())

run = compilation.compile(step)(x_spec, gain=3.0)
out = run(Tensor.ones([4, 2], device=CPU()), gain=3.0)  # "batch" accepts 4

Transforms

as_subgraph	Lowers fn to one shared subgraph body per distinct stage.
compile	Traces and compiles fn.
stage	Traces fn into a graph, without compiling it.

Results

CompiledCallable	A compiled function over tensors.
StagedGraph	A traced graph, ready to inspect or compile.