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

module_dataclass

`module_dataclass()`

max.experimental.nn.module_dataclass(cls=None, /, *, repr=False, **kwargs)

source

Applies dataclass() to a Module subclass.

The decorated class gets an __init__ generated from its field annotations and keeps the Module repr. Parameter tracking does not depend on this decorator.

from max.driver import CPU
from max.dtype import DType
from max.experimental import functional as F
from max.experimental import random
from max.experimental.nn import Module, Linear, module_dataclass
from max.experimental.tensor import Tensor, default_device
from max.graph import TensorType

@module_dataclass
class MLP(Module):
    fc1: Linear
    fc2: Linear

    def forward(self, x: Tensor) -> Tensor:
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        return x

with default_device(CPU()):
    mlp = MLP(
        fc1=Linear(128, 256),
        fc2=Linear(256, 128)
    )

    print(dict(mlp.parameters).keys())

    random.set_seed(0)
    x = random.normal([4, 128], dtype=DType.float32)
    output = mlp(x)

    input_type = TensorType(DType.float32, ["batch", 128], device=mlp.device)
    compiled = mlp.compile(input_type)
    compiled_output = compiled(x)

Parameters:

cls (type[Module[..., Any]] | None) – The class to decorate. Must define a forward method. When None, returns a decorator function (supports using @module_dataclass with or without parentheses).
repr (bool) – If True, use dataclass’s default __repr__ instead of Module’s rich representation. Defaults to False.
**kwargs – Additional keyword arguments forwarded to Python’s @dataclass decorator (e.g., frozen, eq).

Returns:

The decorated class as a Module subclass with automatic parameter tracking and graph compilation capabilities. When cls is None, returns a decorator function.

module_dataclass()​

`module_dataclass()`