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

as_subgraph

as_subgraph()

max.experimental.nn.as_subgraph(fn, *, name=None, prefix=None)

source

Traces a function into a subgraph once and calls it from every call site.

Use this for a block of computation that repeats, such as a transformer layer, so that the compiler processes the definition once instead of once per repetition. The subgraph’s input types are taken from the arguments of the first call, so there is nothing to declare besides fn. Outside a trace, or when the graph was staged with allow_subgraphs=False, the returned function simply calls fn.

Usable as a decorator or at the call site. The following example shares one subgraph between three calls:

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

@compilation.as_subgraph
def block(x: Tensor) -> Tensor:
    return x * 2

x_type = TensorLayout(DType.float32, [4], CPU())
staged = compilation.stage(lambda x: block(block(block(x))))(x_type)

Each distinct callable gets its own subgraph, so two closures over different values are traced separately. The layers of a model are distinct objects that differ only in their weights. To share one subgraph between them, give every call site the same name; each call then resolves its own weights under its prefix, which defaults to the layer’s position in the model:

# ``layers.<i>.<...>``, as ``Module.compile`` names them.
def transformer(x: Tensor, layers: list[Module]) -> Tensor:
    for layer in layers:
        x = compilation.as_subgraph(layer, name="layer")(x)
    return x

Parameters:

• fn (Callable[[~_P], _R]) – The function to trace.
• name (str | None) – The name of the subgraph. Defaults to the name of fn.
• prefix (str | None) – The name prefix that weights used inside fn are relative to. Defaults to the position of fn in its model, read from the names of the weights it holds, or "" if it holds none.

Returns:

A function with the signature of fn that emits one call to the subgraph each time it is called.

Return type:

Callable[[~_P], _R]