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 class

StagedGraph

StagedGraph

class max.experimental.compilation.StagedGraph(fn, graph, layouts, treedef, signal_device_ids=(), *, prefix='', subgraph_cache=None, signal_buffer_index=None)

source

Bases: Generic[_P, _R]

A traced graph that is ready to inspect or compile.

stage() returns one. Print it to see the MLIR of the whole module, including any subgraphs. Pass it to CompiledCallable to compile it.

The following example stages a function and checks its MLIR:

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 scale(x: Tensor) -> Tensor:
    return x * 2

x_type = TensorLayout(DType.float32, [4], CPU())
staged = compilation.stage(scale)(x_type)
print(staged)

Parameters:

• fn (Callable[..., Any]) – The function to trace.
• graph (Graph) – The graph to trace into. Its inputs must be one per shard of each layout, in order, followed by the signal buffers if any.
• layouts (Sequence[TensorLayout]) – The layout of each tensor argument, in graph order.
• treedef (tree.TreeDef) – The tree structure of the arguments, as returned by _declared_layouts.
• signal_device_ids (tuple[int, ...]) – The ids of the accelerators whose collectives need signal buffers.
• prefix (str) – The name prefix that weights created inside fn are relative to.
• subgraph_cache (dict[Any, Any] | None) – The table that as_subgraph() stores traced bodies in. Pass None to inline every subgraph call.
• signal_buffer_index (int | None)

graph

graph: Graph

source

The traced graph.

in_layouts

in_layouts: list[TensorLayout]

source

The layout of each tensor argument, in graph order.

in_tree

in_tree: TreeDef

source

The tree structure of the arguments, including static arguments.

out_layouts

out_layouts: list[TensorLayout]

source

The layout of each tensor result, in graph order.

out_tree

out_tree: TreeDef

source

The tree structure of the return value.