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.compilation.as_subgraph(fn, *, name=None, prefix='', key=<inferred>)

source

Lowers fn to one shared subgraph body per distinct stage.

Usable as a decorator or at the call site.

from max.dtype import DType
from max.experimental import compilation
from max.experimental.tensor import Tensor
from max.graph import DeviceRef, TensorType

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

spec = TensorType(DType.float32, [4], device=DeviceRef.CPU())
staged = compilation.stage(lambda x: block(block(block(x))))(spec)

A shared body also shares the weights it declares. At the call site, prefix gives each site its own weights out of the one body:

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

w_type = TensorLayout(DType.float32, [1], CPU())

def block(x: Tensor) -> Tensor:
    return x * F.constant_external("w", w_type, is_placeholder=True)

def model(x: Tensor) -> Tensor:
    for layer in ("layers.0.", "layers.1."):
        x = compilation.as_subgraph(block, prefix=layer)(x)
    return x

one = Tensor.ones([1], device=CPU())
weights = {"layers.0.w": one * 2, "layers.1.w": one * 10}

run = compilation.compile(model, weights=weights)(w_type)
out = run(one)  # [20.0]

Parameters:

• fn (Callable[[~_P], _R]) – The callable to lower.
• name (str | None) – The subgraph’s name. Defaults to fn’s own name.
• prefix (str) – Prepended to the relative weight names the body declares, so each call site resolves its own weights from a shared body.
• key (str | None) – What identifies this body beyond its arguments, completed here with the argument structure and operand types. Pass None to compare the staged IR instead. Omitted, a key is derived from fn where that is sound.

Returns:

A callable with fn’s signature that emits a call to the shared body.

Raises:

TypeError – If called outside a capture. Call fn directly to run eagerly.

Return type:

Callable[[~_P], _R]