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

Module

Module

class max.nn.Module

source

Bases: Layer, ABC

Base class for model components with weight management.

Provides functionality to create custom layers and construct networks with automatic weight tracking.

The following example uses the Module class to create custom layers and build a neural network:

from max import nn
from max.dtype import DType
from max.graph import Weight, ops, DeviceRef

class Linear(nn.Module):
    def __init__(self, in_dims, out_dims):
        super().__init__()
        self.weight = Weight("weight", DType.float32, (in_dims, out_dims), DeviceRef.CPU())

    def __call__(self, x):
        return x @ self.weight.T

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.up = Linear(5, 10)
        self.gate = Linear(5, 10)
        self.down = Linear(10, 5)

    def __call__(self, x):
        return self.down(ops.silu(self.gate(x)) + self.up(x))

model = MLP()
print(model.state_dict())  # {"up.weight": Buffer([5, 10]), ...}

Constructing a graph without Module can result in name collisions with the weights (in this example, there would be three weights with the name Weight). With Module, you can use state_dict() or load_state_dict() to initialize or set the weights values, and finalize the weight names to be unique within the model.

build_subgraph()

build_subgraph(name, inputs, weight_prefix='')

source

Builds a subgraph encapsulating this layer’s computation.

Call this method once on a representative layer, then call the returned subgraph once per layer using call() with a unique prefix. This pattern lets the compiler process the layer definition once rather than once per repetition, which significantly reduces compile time for models with many identical layers.

For example, define a small layer, build a subgraph from a representative input, then call it once per repetition with a layer-specific weight prefix:

from max.dtype import DType
from max.graph import DeviceRef, Graph, TensorType, Weight, ops
from max.nn.layer import Module

class Linear(Module):
    def __init__(self, in_dims, out_dims, weight_name):
        super().__init__()
        self.weight = Weight(
            weight_name,
            DType.float32,
            (in_dims, out_dims),
            DeviceRef.CPU(),
        )

    def __call__(self, x):
        return x @ self.weight.T

num_layers = 3
# prefix so build_subgraph() turns it into a placeholder.
layer = Linear(4, 4, weight_name="layers.0.weight")
with Graph(
    "build_subgraph_example",
    input_types=(
        TensorType(DType.float32, [2, 4], device=DeviceRef.CPU()),
    ),
) as graph:
    h = graph.inputs[0].tensor

    # Build the subgraph once from a representative input.
    subgraph = layer.build_subgraph(
        "linear_block",
        inputs=[h],
        weight_prefix="layers.0.",
    )

    # Call it once per layer with the correct weight prefix.
    for idx in range(num_layers):
        outputs = ops.call(subgraph, h, prefix=f"layers.{idx}.")
        h = outputs[0].tensor

    graph.output(h)

Parameters:

• name (str) – The name of the subgraph. Must be unique within the containing graph.
• inputs (Sequence[Tree[Any]]) – Representative input values for the subgraph, one per positional argument of the layer’s __call__. Each argument may be a single Value, a (possibly nested) list/tuple of values, or a structured type following the pytree protocol, such as PagedCacheValues. The subgraph’s signature is derived from the flattened leaves’ types, and the same structure is rebuilt from the subgraph’s inputs before the layer is invoked.
• weight_prefix (str) – A prefix string to strip from weight names before registering them as placeholder weights. At call time, the caller supplies the same prefix via the prefix argument of call() to re-resolve each weight to the correct entry in the weights registry.

Returns:

A Graph instance representing the subgraph.

Return type:

Graph

Notes:

Weights with names that start with weight_prefix are marked as placeholders. Any call() invocation for this subgraph must supply a matching prefix. Pass the flattened leaves (via call()) in the same order that inputs flattens to.

compile()

compile(*input_types, session, weights=None, name=None, custom_extensions=(), is_device_graph=False)

source

Traces this module, then compiles and initializes the graph.

The compile goes through session’s precompiled or exported MEFs when the session has them, as InferenceSession.load does.

Parameters:

• *input_types (Type[Any]) – The type of each graph input.
• session (InferenceSession) – The session that compiles and initializes the graph.
• weights (Mapping[str, DLPackArray] | None) – The weight values, keyed by the graph’s weight names. Defaults to state_dict().
• name (str | None) – The name of the graph. Defaults to the class’s qualified name.
• custom_extensions (Iterable[Path]) – Paths to custom kernel libraries the graph uses.
• is_device_graph (bool) – Whether to build a device graph.

Returns:

The compiled and initialized Model.

Return type:

Model

layer_weights

property layer_weights: dict[str, Weight]

source

Returns a mapping from weight name to Weight for this layer.

load_state_dict()

load_state_dict(state_dict, *, override_quantization_encoding=False, weight_alignment=None, strict=True)

source

Sets the values of all weights in this model.

The keys in state_dict must match the fully-qualified weight names used internally by the Module. Those names normally follow the attribute hierarchy (such as model.layers.0.self_attn.qkv_proj.weight), but a sublayer whose _omit_module_attr_name is True is omitted from its descendants’ FQNs. The canonical example is StackedLinear in unfused mode, where self.qkv_proj = StackedLinear(names=["q_proj", "k_proj", "v_proj"], stacked=False) exposes weights at self_attn.q_proj.weight / self_attn.k_proj.weight / self_attn.v_proj.weight rather than nested under self_attn.qkv_proj.. Use raw_state_dict() to inspect the exact keys this method expects for a given module.

Parameters:

• state_dict (Mapping[str, DLPackArray | WeightData]) – A map from weight name to a numpy array or Buffer.
• override_quantization_encoding (bool) – Whether to override the weight quantization based on the loaded value.
• weight_alignment (int | None) – If specified, overrides the alignment for each weight in the Module. If left as None, each value in state_dict must be aligned by the default dtype alignment.
• strict (bool) – If True, raises an error if any weights required by the Module are missing from state_dict, or if any keys in state_dict were not used by the Module. If False, both missing and unexpected keys are tolerated and reported only via return values/logging by callers.

Raises:

ValueError – If strict is True and any required weight is missing from state_dict, or if state_dict contains keys not used by the Module.

Return type:

None

named_modules()

named_modules()

source

Yields (weight prefix, module) for every module in the tree.

The walk is bfs from self and visits each module once, even when it is reachable through several attribute paths. The prefix is the string raw_state_dict prepends to the module’s own weight names: empty for self, otherwise the attribute path with a trailing . (e.g. layers.3.mlp.). A sublayer whose _omit_module_attr_name is True shares its parent’s prefix.

Returns:

An iterable of (prefix, module) pairs.

Return type:

Iterable[tuple[str, Module]]

raw_state_dict()

raw_state_dict()

source

Returns all weights objects in the model. Unlike state_dict(), this returns Weight objects instead of the assigned values. Some parameters inside the Weight can be configured before a graph is built. Do not change these attributes after building a graph:

• align
• dtype
• quantization_encoding
• shape

Keys follow the same FQN convention as load_state_dict(): attribute paths through the module tree, with any sublayer that sets _omit_module_attr_name skipped in the prefix.

Returns:

Map from weight name to the Weight object.

Return type:

dict[str, Weight]

replace_module()

replace_module(name, module)

source

Replaces the child module registered at name.

Parameters:

• name (str) – The attribute name of an existing child module.
• module (Module) – The replacement module.

Raises:

KeyError – If name is not a registered child module. Unlike a bare setattr, this surfaces a wrong name instead of silently creating a new, unreferenced child.

Return type:

None

set_shared_weight()

set_shared_weight(name, weight)

source

Registers a Weight as shared on this layer.

Sets name as an attribute on this layer and marks the weight as shared so that raw_state_dict() and load_state_dict() skip it when iterating over owned weights.

Parameters:

• name (str) – The attribute name under which the weight is registered.
• weight (Weight) – The Weight to share.

Return type:

None

state_dict()

state_dict(auto_initialize=True)

source

Returns values of all weights in the model.

The values returned are the same as the values set in load_state_dict(). If load_state_dict() has not been called and none of the weights have values, then they are initialized to zero.

Keys follow the same FQN convention as load_state_dict(): attribute paths through the module tree, with any sublayer that sets _omit_module_attr_name (such as StackedLinear in unfused mode) skipped in the prefix.

Parameters:

auto_initialize (bool) – Determines whether to initialize weights to zero if the weight value has not been loaded. If this is False, a ValueError is raised if an uninitialized weight is found.

Returns:

Map from weight name to the weight value (can be numpy array or Buffer).

Return type:

dict[str, DLPackArray]

sublayers

property sublayers: dict[str, Module]

source

trace()

trace(*input_types, name=None, module=None, custom_extensions=(), is_device_graph=False)

source

Traces this module’s __call__ into a graph without compiling it.

The graph passes one input per entry of input_types to __call__ positionally, and __call__ returns the graph’s outputs: one value, a sequence of values, or None. A subclass whose __call__ takes structured arguments overrides this method to adapt the flat graph inputs to them. Weights keep the names they have when tracing starts. load_state_dict() and state_dict() name each weight by its attribute path, so call one of them first when the graph’s weights registry comes from state_dict().

The following example traces a linear layer:

from max.dtype import DType
from max.graph import DeviceRef, TensorType, Weight
from max.nn.layer import Module

class Linear(Module):
    def __init__(self) -> None:
        super().__init__()
        self.weight = Weight(
            "weight", DType.float32, (4, 4), DeviceRef.CPU()
        )

    def __call__(self, x):
        return x @ self.weight.T

graph = Linear().trace(
    TensorType(DType.float32, [2, 4], device=DeviceRef.CPU())
)

Parameters:

• *input_types (Type[Any]) – The type of each graph input.
• name (str | None) – The name of the graph. Defaults to the class’s qualified name.
• module (Module | None) – The MLIR module to add the graph to, so that several graphs compile into one artifact. Defaults to a new module.
• custom_extensions (Iterable[Path]) – Paths to custom kernel libraries the graph uses.
• is_device_graph (bool) – Whether to build a device graph.

Returns:

The traced Graph.

Return type:

Graph