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

Traces and compiles Python functions over tensors.

compile() takes a function that operates on Tensor values and returns a CompiledCallable that runs it as a compiled graph. stage() performs the tracing step only and returns a StagedGraph, which is useful for inspecting the generated MLIR.

Both take the function first and its input types second. An input type describes one tensor argument: pass a TensorType or a TensorLayout for an argument the function reads, or a BufferType or BufferLayout for one it also writes to. Arguments that are not tensors, such as a Python float or bool, are static: their values are fixed into the graph during tracing, so pass the same values when calling the compiled function.

The following example compiles a function with one tensor argument and one static argument:

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_type = TensorLayout(DType.float32, ["batch", 2], CPU())

run = compilation.compile(step)(x_type, gain=3.0)
out = run(Tensor.ones([4, 2], device=CPU()), gain=3.0)

Transforms

`as_subgraph`	Traces a function into a subgraph once and calls it from every call site.
`compile`	Traces and compiles a function.
`stage`	Traces a function into a graph without compiling it.

Results

`CompiledCallable`	A compiled function that runs on tensors.
`StagedGraph`	A traced graph that is ready to inspect or compile.

Transforms​

Results​

Transforms

Results