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

source

Bases: Generic[_P, _R]

A traced graph, ready to inspect or compile.

Returned by stage(). Printing one renders the MLIR of the whole module, subgraph bodies included. Call compile() to make it runnable.

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

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

compile()

compile(*, weights=None)

source

Compiles the graph into a CompiledCallable.

Parameters:

weights (Mapping[str, DLPackArray] | None) – Data for the external constants the graph declares, keyed as the graph names them, one entry per shard of a distributed weight.

Returns:

The compiled function, called on real tensors.

Return type:

CompiledCallable[~_P, _R]

graph

graph: Graph

source

The recorded graph.