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

Linear

Linear

class max.experimental.nn.Linear(in_dim, out_dim, *, bias=True, quant_config=None)

source

Bases: Module

A unary linear transformation over an input tensor.

Linear is defined as f(x) = x @ W.T + B where W is the weight tensor and B is an optional bias tensor.

If W is not square then the transformation represents a dimensionality change. By convention the weight tensor is stored transposed.

from max.experimental.nn import Linear
from max.experimental.tensor import Tensor

model = Linear(5, 10)

assert dict(model.parameters) == {
    "weight": model.weight, "bias": model.bias
}

result = model(Tensor.ones([5]))
assert result.shape == [10]

Constructs a random linear transformation of the given dimensions.

Parameters:

• in_dim (DimLike) – The dimensionality of the input to the transformation
• out_dim (DimLike) – The dimensionality after applying the transformation to the input tensor of dim in_dim.
• bias (Tensor | Literal[0]) – Whether to use a bias in the transformation.
• quant_config (QuantConfig | None) – The QuantConfig of a quantized weight. Only static per-tensor FP8 is supported: an float8_e4m3fn weight with a weight_scale and an input_scale, which start as zeros.

Raises:

NotImplementedError – If quant_config is not static per-tensor FP8.

bias

bias: Tensor | Literal[0]

source

The bias Tensor for the linear transformation (or 0 if bias is disabled).

forward()

forward(x)

source

Applies a linear transformation to the input tensor.

Linear is defined as f(x) = x @ W.T + B where W is the weight tensor and B is an optional bias tensor.

Parameters:

x (Tensor) – The input tensor

Returns:

The result of applying the linear transformation to the tensor.

Return type:

Tensor

in_dim

property in_dim: Dim

source

The input dimension for the transformation.

input_scale

input_scale: Annotated[Tensor, _DevicePinned]

source

The static per-tensor input scale, on the CPU. Present only with a quant_config.

out_dim

property out_dim: Dim

source

The output dimension for the transformation.

weight

weight: Tensor

source

The weight Tensor for the linear transformation.

weight_scale

weight_scale: Annotated[Tensor, _DevicePinned]

source

The per-tensor weight scale, on the CPU. Present only with a quant_config.