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.experimental.nn.Module

source

Bases: Generic[_P, _R]

The core unit of composition for modeling in MAX.

Informally, a Module is a container class. It can contain other Module instances, tensors (the Module’s “local parameters”) or other arbitrary Python data.

A Module also has a forward() method which defines how the Module computes its output. In the simplest case this is a function from one tensor to another tensor. Users call the module using __call__() which internally invokes forward().

Formally modules form a tree, and subtrees of modules can be manipulated directly. A Module may also be thought of as a closure, where the parameters form the data of the closure and forward() is the application of the closure.

Users who do not use a Python type checker, or use lax settings for their type checker, may inherit from Module without parameters. Users who use a type checker with stricter settings (including MAX internal code) should specify explicit types for full type checking:

class Linear(Module[[Tensor], Tensor]):
    def forward(self, x: Tensor) -> Tensor:
        return x @ self.weight.T + self.bias

Terminology:

• A “child” of a Module is a sub-Module stored directly on that Module.
• A “descendant” of a Module is one of its children, or one of their descendants.
• A “parameter” is a tensor storing data on the Module or one of its descendants.
• The “qualified path” of a descendant is a period-separated string of the names of the child module attributes which lead to that descendant module, for instance child.sub.last.
• The “qualified path” of a parameter is the qualified path of the descendant directly holding that parameter, followed by a final path component for the attribute name of the tensor. For instance weight for a local parameter, or child.sub.last.weight for a descendant’s parameter.

from max.experimental.tensor import Tensor
from max.experimental.nn import Module, module_dataclass

@module_dataclass
class Linear(Module):
    weight: Tensor
    bias: Tensor | int = 0

    def forward(self, x: Tensor) -> Tensor:
        return x @ self.weight.T + self.bias

linear = Linear(Tensor.zeros([5, 4]))
print(linear)
print(linear(Tensor([1, 2, 3, 4])))

Device placement:

MAX uses a compiled graph model that separates weight storage from computation placement. Understanding this distinction is essential for running models on GPU.

to() is the single pre-compilation entry point for device placement. It moves all weight tensors to the target device and records it on the module via the device property. When you construct the TensorType objects you pass to compile(), reference model.device for their device field, so a single to() call drives both weight placement and computation placement:

from max.dtype import DType
from max.experimental.nn import Linear
from max.experimental.tensor import Tensor, defaults
from max.graph import TensorType

model = Linear(5, 10)
_, mesh = defaults()
model.to(mesh.devices[0])

input_type = TensorType(DType.float32, ["batch", 5], device=model.device)
compiled = model.compile(input_type)
result = compiled(Tensor.ones([3, 5], dtype=DType.float32))

For CPU (the default), calling to() is optional. The device property defaults to CPU:

from max.driver import CPU
from max.experimental.nn import Linear

model = Linear(5, 10)
print(model.device)

Because device is tracked per-module instance, sub-modules can be placed on different devices independently. Here two Linear sub-modules are placed on separate CPU device references (use distinct Accelerator instances when accelerators are available):

from max.driver import CPU
from max.experimental.nn import Linear

encoder = Linear(5, 8)
decoder = Linear(8, 4)

encoder.to(CPU(0))
decoder.to(CPU(0))

For graph-level tensor routing inside forward() (such as pulling an activation back to CPU at the end of the graph), use transfer_to() or to() instead; those insert transfer nodes into the compiled graph and are unrelated to pre-compilation weight placement.

API	When it runs	What it moves
Module.to(device)	Python host, before compile()	Stored weight tensors; records module.device
ops.transfer_to(x, d) / TensorValue.to(d)	Graph execution time (inside forward())	Activation tensors within the compiled graph
Tensor.to(device)	Eager runtime (outside a graph)	Concrete eager tensors (such as staging inputs)

children

property children: Iterable[tuple[str, Module[..., Any]]]

source

Iterates over the direct child modules of the Module.

Yields:

(name, module) pairs, where name is the attribute name of the child on the module.

compile()

compile(*input_types, weights=None, session=None, name=None, custom_extensions=(), auto_cast=False, allow_subgraphs=True, is_device_graph=False, signal_buffer_index=None)

source

Traces forward and compiles it into a callable.

The graph refers to each parameter by name, and weights supplies the values. Call to() before compiling and build the input types from device, so that the parameters and the computation are on the same device.

The following example compiles a linear layer and runs it:

from max.dtype import DType
from max.experimental.nn import Module, module_dataclass
from max.experimental.tensor import Tensor, TensorType, defaults

@module_dataclass
class Linear(Module):
    weight: Tensor
    bias: Tensor

    def forward(self, x: Tensor) -> Tensor:
        return x @ self.weight.T + self.bias

linear = Linear(
    Tensor.zeros([10, 5], dtype=DType.float32),
    Tensor.zeros([10], dtype=DType.float32),
)
_, mesh = defaults()
run = linear.compile(
    TensorType(DType.float32, [3, 5], mesh.devices[0])
)
result = run(Tensor.ones([3, 5], dtype=DType.float32))

Parameters:

• *input_types (Type[Any] | TensorLayout) – The type of each tensor argument of forward, as a TensorType or a TensorLayout. The device of each type determines where the computation runs.
• weights (Mapping[str, DLPackArray] | None) – The values of the parameters, keyed by the names that parameters yields. Defaults to the values currently held by this module.
• session (InferenceSession | None) – The session that compiles and initializes the graph, including its precompiled or exported MEFs and its kernel options. Defaults to a session shared by the whole process.
• name (str | None) – The name of the graph. Defaults to the class’s qualified name.
• custom_extensions (Iterable[Path]) – Paths to custom Mojo kernel libraries, required when forward calls custom().
• auto_cast (bool) – Whether to permit a safe float32/bfloat16 cast when shapes match. See load_state_dict().
• allow_subgraphs (bool) – Whether calls made through as_subgraph() are traced as subgraphs. If False, they are inlined.
• is_device_graph (bool) – Whether to build a device graph.
• signal_buffer_index (int | None) – The graph input position of the signal buffers. See stage().

Returns:

A CompiledCallable with the signature of forward.

Raises:

• TypeError – If an input type does not match forward’s signature, or an operation in forward cannot be traced.
• ValueError – If a weight does not fit the parameter it loads into.
• KeyError – If weights has no entry for a parameter.

Return type:

CompiledCallable[~_P, _R]

descendants

property descendants: Iterable[tuple[str, Module[..., Any]]]

source

Iterates over the Module’s descendant modules.

Yields:

(name, module) pairs, where name is the qualified path of the descendant with respect to the module.

device

property device: Device

source

The canonical device for this module’s weights and computation.

Set by calling to() or by assigning self.device in a subclass __init__. When neither has been called the property returns CPU as a safe default so that modules without an explicit device placement still compile and run on CPU. When constructing the TensorType objects you pass to compile(), reference self.device for their device field so that a single to() call drives both weight placement and computation placement.

Returns:

The device this module is placed on, defaulting to CPU if to() has not been called and self.device has not been set in a subclass __init__.

forward()

forward(*args, **kwargs)

source

Defines the computation performed by the module.

Users must override this method in their subclass to define the module’s computation.

Parameters:

• *args (~_P) – Positional arguments for the computation.
• **kwargs (~_P) – Keyword arguments for the computation.

Returns:

The result of applying the module to the input.

Raises:

NotImplementedError – If the subclass does not override this method.

Return type:

_R

load_state_dict()

load_state_dict(state, strict=True, *, auto_cast=False)

source

Loads parameter values from a dictionary into the module hierarchy.

This method updates all module parameters in-place by loading values from the provided state dictionary. The dictionary maps qualified parameter names (dot-separated paths like fc1.weight) to tensor values.

The strict mode (default) ensures all weights in the dictionary are actually used, catching errors from mismatched architectures or incorrect weight names.

For example, the following loads weights from a dictionary into a model:

from max.experimental.tensor import Tensor
from max.experimental.nn import Module, module_dataclass

@module_dataclass
class Linear(Module):
    weight: Tensor
    bias: Tensor

    def forward(self, x: Tensor) -> Tensor:
        return x @ self.weight.T + self.bias

model = Linear(
    weight=Tensor.zeros([10, 5]),
    bias=Tensor.zeros([10])
)

weights = {
    "weight": Tensor.zeros([10, 5]),
    "bias": Tensor.zeros([10]),
}
model.load_state_dict(weights)

By default (auto_cast=False) any dtype mismatch between the loaded tensor and the parameter raises. When auto_cast=True, loaded weights whose dtype is in the safe-cast set (currently float32 and bfloat16) are automatically cast to the parameter’s dtype when shapes match, and a single summary message is logged at WARNING level per call describing how many parameters were cast. Narrowing casts (such as float32 -> bfloat16) are flagged in the log message as (precision loss). Dtype mismatches outside the safe-cast set still raise regardless of auto_cast.

Parameters:

• state (Mapping[str, DLPackArray]) – Dictionary mapping qualified parameter names to tensor values. Keys should match the names from Module.parameters property. Values should be DLPack-compatible arrays or Tensor objects. Shapes must match the existing parameters with the corresponding name. Dtypes must match exactly or both lie in the safe-cast set above. Values may be on a different device; in that case the tensor is copied to the existing parameter’s device.
• strict (bool) – If True (default), verify that all keys in state are used (that is, match actual parameters). If False, silently ignore extra keys that don’t match any parameters.
• auto_cast (bool) – If True, permit safe dtype auto-casting between float32 and bfloat16 when shapes match. Defaults to False — dtype mismatches always raise. Pipelines that want to opt in via MODULAR_AUTO_CAST_WEIGHTS should pass auto_cast=max.pipelines.lib.weight_loading.auto_cast_weights_from_env().

Raises:

• ValueError – If strict=True and some weights in state don’t match any model parameters (indicates architecture mismatch or incorrect weight names).
• ValueError – If a loaded tensor has a different shape than the existing parameter, or a dtype mismatch that is not covered by the safe-cast set (or auto_cast=False).
• KeyError – If a required parameter name in the model is missing from state (regardless of strict setting).

Return type:

None

local_parameters

property local_parameters: Iterable[tuple[str, Tensor]]

source

Iterates over the tensors held directly on this Module.

Yields:

(name, tensor) pairs, where name is the attribute name of the tensor on the module.

name_transparent

name_transparent: bool = False

source

Whether this module’s own name is dropped from its descendants’ paths.

parameters

property parameters: Iterable[tuple[str, Tensor]]

source

Iterates over all parameters in this module and its sub-modules.

This property performs a depth-first traversal of the module hierarchy, yielding each parameter tensor with its qualified name. The qualified name uses dot-notation to represent the module tree structure (such as encoder.layer1.weight).

Parameters are yielded in depth-first order: first the current module’s direct parameters, then recursively each sub-module’s parameters.

Yields:

(name, parameter) tuples where name is the dot-separated qualified path of the parameter and parameter is the Tensor.

stage()

stage(*input_types, weights=None, name=None, custom_extensions=(), auto_cast=False, allow_subgraphs=True, is_device_graph=False, signal_buffer_index=None)

source

Traces forward and builds the weight registry it binds.

This is the tracing half of compile(). Use it to compile and initialize the graph yourself, then wrap the initialized model with from_engine_model().

Parameters:

• *input_types (Type[Any] | TensorLayout) – As compile().
• weights (Mapping[str, DLPackArray] | None) – As compile().
• name (str | None) – As compile().
• custom_extensions (Iterable[Path]) – As compile().
• auto_cast (bool) – As compile().
• allow_subgraphs (bool) – As compile().
• is_device_graph (bool) – As compile().
• signal_buffer_index (int | None) – As compile().

Returns:

The StagedGraph and the weight registry: a mapping from the name of each weight the graph declares to its data, with one entry per shard for a distributed weight.

Raises:

• TypeError – As compile().
• ValueError – As compile().
• KeyError – As compile().

Return type:

tuple[StagedGraph[~_P, _R], dict[str, Tensor]]

to()

to(target)

source

Transfers all module parameters to a single device.

Records the target as device. Build the input types you pass to compile() from it, so that the computation runs where the parameters are.

Parameters:

target (Device | DeviceMesh | DeviceMapping) – The device for all module parameters.

Returns:

A reference to the model. The transfer is applied mutably; the module’s device property and all internal parameters are updated in place.

Raises:

ValueError – If target spans more than one device.

Return type:

Self

trace()

trace(*input_types, name=None, custom_extensions=(), allow_subgraphs=True, is_device_graph=False, signal_buffer_index=None)

source

Traces forward into a graph without compiling it.

The graph refers to each parameter by name, as an external constant, rather than embedding its value. compile() supplies the values from the weight registry. Use this method to inspect the MLIR that compile() would compile.

The following example traces a linear layer and prints its MLIR:

from max.driver import CPU
from max.dtype import DType
from max.experimental.nn import Linear
from max.experimental.tensor import TensorType, default_device

with default_device(CPU()):
    linear = Linear(4, 8)
    x_type = TensorType(DType.float32, [2, 4], device=linear.device)
    staged = linear.trace(x_type)
    print(staged)

Parameters:

• *input_types (Type[Any] | TensorLayout) – The type of each tensor argument of forward, as a TensorType or a TensorLayout.
• name (str | None) – The name of the graph. Defaults to the class’s qualified name.
• custom_extensions (Iterable[Path]) – Paths to custom Mojo kernel libraries.
• allow_subgraphs (bool) – Whether calls made through as_subgraph() are traced as subgraphs. If False, they are inlined.
• is_device_graph (bool) – Whether to build a device graph.
• signal_buffer_index (int | None) – The graph input position of the signal buffers. See stage().

Returns:

The StagedGraph. Print it to see the MLIR.

Return type:

StagedGraph[…, Any]