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

subgraphable

`subgraphable()`

max.experimental.nn.subgraphable(module, *, name=None)

source

Marks a Module to call to a subgraph.

Use this to speed up graph build and compilation time.

When a name is given, all modules with that name will share the same subgraph. When a name is not given, the IR of the module’s traced call is used to determine whether an existing subgraph can be used (this process is slow, but guarantees correctness).

Warning

Kernel fusion cannot cross a subgraph boundary, so prefer marking larger modules, like encoder/decoder blocks.

Use it as a class decorator so an ordinary layer loop auto-shares a body:

@subgraphable
@module_dataclass
class Block(Module[[Tensor], Tensor]): ...

def forward(self, x):
    for layer in self.layers:  # each call reuses one shared subgraph
        x = layer(x)
    return x

Or mark a single instance directly: subgraphable(layer, name="block").

Parameters:

module (_T) – The Module subclass (class-decorator form) or instance to mark.
name (str | None) – Optional subgraph key. Modules marked with the same name share one subgraph definition; omit it to deduplicate subgraphs based on the module’s traced call instead.

Returns:

The same class or instance, marked so Module.__call__() lowers each call into a shared subgraph.

Raises:

TypeError – If module is not a Module subclass or instance.

Return type:

subgraphable()​

`subgraphable()`