Module layers 

Layers, as in MLX’s mlx.nn (and mlx-lm’s LoRA layer). Initialisation draws from this thread’s key sequence in MLX’s order, so the same seed gives MLX’s values.

Constructors take no bool and return no Result: they build lazy arrays, and an invalid size gives failed parameters whose error surfaces at the first boundary. Every forward returns an Array and is #[track_caller]: an invalid input gives a failed array whose error (check(), eval(), …) reports the line that called forward.

Structs

Dropout

Randomly zeroes elements with probability p during training and scales the rest by 1/(1 − p); the identity in evaluation mode (MLX’s nn.Dropout). The mask is bernoulli(1 − p), drawn from the thread’s key sequence (bit for bit MLX’s for a seed) or from an explicit key (Dropout::forward_with_key).

Embedding

A lookup table of dims-dimensional vectors for integer ids (MLX’s nn.Embedding).

LayerNorm

Layer normalisation over the last axis, optionally with a learned scale and bias (MLX’s nn.LayerNorm, the fused fast::layer_norm).

Linear

An affine layer y = x Wᵀ + b, with W of shape [output, input] and b of shape [output] (MLX’s nn.Linear).

LoraLinear

A frozen linear layer plus a trainable low-rank update (mlx-lm’s LoRALinear): y = base(x) + scale · (dropout(x) A) B, with A of shape [input, r] and B of shape [r, output] (zeros at first, so the layer starts as its base). Names follow mlx-lm: linear.weight, lora_a, lora_b.

QuantizedEmbedding

An Embedding with its table quantized (MLX’s nn.QuantizedEmbedding): rows are gathered packed and dequantized; the packed table, scales and biases are buffers. In the floating-point modes the rows come out in bfloat16, as MLX’s.

QuantizedLinear

A Linear layer with its weight quantized (MLX’s nn.QuantizedLinear): x Wᵀ + b computed by quantized_matmul from the packed weight and its scales (and, for affine, biases), in any QuantMode.

RmsNorm

Root-mean-square normalisation over the last axis with a learned scale (MLX’s nn.RMSNorm, the fused fast::rms_norm).

Rope

Rotary position embedding of the first dims features (MLX’s nn.RoPE, the fused fast::rope), optionally with a frequency scaling (Rope::scaled: mlx-lm’s Llama3RoPE, YarnRoPE, SuScaledRoPE, ProportionalRoPE). No parameters: the frequencies are computed from the settings, never loaded or saved.

Enums

LoraBase

The layer a LoraLinear adapts.