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 module

max.experimental.nn

Module framework for max.experimental.

Provides the Module base class, the module_dataclass() decorator, and Module.compile(), which traces forward against symbolic inputs and returns an ahead-of-time-compiled CompiledCallable. The same forward body runs eagerly, traces to a Graph, or compiles, depending on the active realization context.

Example:

from max.driver import CPU
from max.dtype import DType
from max.experimental.tensor import Tensor, default_device
from max.experimental.nn import Module, module_dataclass
from max.graph import TensorType

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

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

with default_device(CPU()):
    model = MyLayer(weight=Tensor.zeros([10, 5]), bias=Tensor.zeros([10]))
    y = model(Tensor.ones([3, 5]))

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

Submodules

• max.experimental.nn.common_layers.functional_kernels
• max.experimental.nn.norm
• max.experimental.nn.rope

Base classes

Module	The core unit of composition for modeling in MAX.
ModuleList	A Module subclass which is locally a list container.
PinnedDeviceTensor	Runtime representation of an annotated type.
Sequential	A Module subclass which holds a sequence of unary modules.

Linear layers

Conv2d	A 2D convolution layer.
Embedding	A vector embedding.
Linear	A unary linear transformation over an input tensor.

Hyper-connection

HyperConnection	Mixes hc_mult parallel residual streams around a transformer sublayer.

Functions

as_subgraph	Traces a function into a subgraph once and calls it from every call site.
module_dataclass	Applies dataclass() to a Module subclass.