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

Module framework for models built on max.tensor.

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.tensor import Tensor, default_device
from max.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.nn.legacy	Neural network modules for MAX.
max.nn.norm	Normalization layers for MAX neural networks.
max.nn.rope	Positional embedding modules and functions.

Base classes

Module	The core unit of composition for modeling in MAX.
ModuleList	A Module subclass which is locally a list container.
PinnedDeviceTensor	Type alias for a Tensor parameter that Module.to will leave on its current device.
Sequential	A Module subclass which holds a sequence of unary modules.
TransparentModule	A Module that drops its own name from its descendants' paths.

Linear layers

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

Attention and rotary embeddings

AttentionWithRope	Implementation of attention that uses Rotary Position Embedding (RoPE).
RotaryEmbedding	RotaryEmbedding layer to calculate and apply the frequency tensor for complex exponentials.
YarnRotaryEmbedding	Generic YaRN (Yet another RoPE eNhancement) Rotary Position Embedding layer.

LoRA

LoRA	Wraps a projection module with an additive multi-adapter LoRA term.

lora_layers	Yields (qualified_name, layer) for every LoRA in model.
lora_parameters	Yields (qualified_name, slot) for every adapter stack in model.

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.

Mesh axes

DP	Mesh axis name for data parallelism.
EP	Mesh axis name for expert parallelism.
TP	Mesh axis name for tensor parallelism.