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.tensor.compilation.stage(fn, *, name=None, custom_extensions=(), allow_subgraphs=True, signal_devices=(), is_device_graph=False, signal_buffer_index=None)

source

Traces a function into a graph without compiling it.

Returns a function that takes the input types of fn and returns a StagedGraph. Use it to inspect the MLIR that compile() would compile.

The following example stages a function whose tensor arguments are passed in a dictionary:

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

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

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

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

Parameters:

fn (Callable[[~_P], _R]) – The function to trace. It takes and returns Tensor values, possibly inside lists, tuples or dictionaries.
name (str | None) – The name of the graph. Defaults to the name of fn.
custom_extensions (Iterable[Path]) – Paths to custom Mojo kernel libraries.
allow_subgraphs (bool) – Whether calls made through as_subgraph() are traced as subgraphs. If False, they are inlined.
signal_devices (Iterable[Device]) – Additional accelerators to allocate signal buffers for, beyond the devices of the input types and of the tensors held by fn.
is_device_graph (bool) – Whether to build a device graph.
signal_buffer_index (int | None) – The graph input position of the signal buffers, counted over the flattened input types. Defaults to after the last input. Set it to match a calling convention that fixes their position, such as an exported graph’s. Setting it keeps the signal buffer on a single accelerator, which otherwise needs none.

Returns:

A function that takes one input type per tensor argument of fn and returns the StagedGraph.

Return type:

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

stage()​

`stage()`