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 function

stage

`stage()`

max.experimental.compilation.stage(fn, *, name=None, custom_extensions=(), allow_subgraphs=True, signal_devices=(), is_device_graph=False)

source

Traces fn into a graph, without compiling it.

Call the returned function with one spec per tensor argument of fn to get the StagedGraph, which prints as MLIR. Use compile() to run fn instead.

Tensor arguments are given as specs. A TensorLayout is a layout, and an argument given one is read-only; pass a BufferLayout for a buffer the callable may also store through. Only a layout may declare a boundary: a live Tensor is refused, because its dims are whatever it currently holds and taking them would fix every one. Pass tensor.layout to do that on purpose.

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

def combine(kv: dict[str, Tensor], alpha: float) -> Tensor:
    return (kv["a"] + kv["b"]) * alpha

spec = TensorLayout(DType.float32, [2], CPU())

staged = compilation.stage(combine)({"a": spec, "b": spec}, 2.0)
print(staged)

Parameters:

fn (Callable[[~_P], _R]) – The callable to record, over Tensor values or containers of them.
name (str | None) – The graph’s name. Defaults to fn’s own name.
custom_extensions (Iterable[Path]) – Paths to custom Mojo kernel libraries.
allow_subgraphs (bool) – Whether as_subgraph() bodies become shared subgraphs rather than inlining into the caller.
signal_devices (Iterable[Device]) – Devices taking part in collectives beyond what the specs span.
is_device_graph (bool) – Whether to record a device graph.

Returns:

A callable taking one spec per argument of fn, as as_layout() accepts them, and returning the StagedGraph.

Return type:

Callable[[…], StagedGraph[~_P, _R]]

stage()​

`stage()`