Crate nuvec_nn 

Neural-network building blocks on nuvec, following MLX’s mlx.nn and mlx.optimizers.

A model is a plain struct of layers with #[derive(Module)]: its parameters get MLX’s names from the field names (layers.0.weight), and every model is a Tree, so the core transforms differentiate it and return gradients of the model’s own type (grads.layers[0].weight).

• Module: names and counts (Module::parameters, Module::num_params, Module::summary, Module::display), freezing by path (Module::freeze, Module::unfreeze), train/eval mode (Module::train), loading and saving by name (Module::load_weights, Module::save_weights) and replacing submodules (Module::map_modules).
• Layers (layers): Linear, QuantizedLinear, Embedding, QuantizedEmbedding, RmsNorm, LayerNorm, Rope, LoraLinear and Dropout; activations relu, silu, gelu_approx and gelu_fast_approx. Constructors take no bool and return no Result: Linear::new(784, 32), Linear::new(32, 10).without_bias().
• losses::cross_entropy and optim: SGD, Adam, AdamW, Lion and Muon with learning-rate schedules, parameter groups (optim::MultiOptimizer) and gradient clipping; the optimizer state is a tree keyed by parameter path.

Layers’ forwards, activations and losses return an Array, as every nuvec op: an invalid input gives a failed array whose error surfaces at a boundary (check(), eval(), item(), a transform) and reports where forward was called.

use nuvec::{Array, Result, array, random, value_and_grad};
use nuvec_nn::optim::{Optimizer, Sgd};
use nuvec_nn::{Linear, Module, Tree, losses, relu};

#[derive(Clone, Module)]
struct Mlp {
    hidden: Linear,
    out: Linear,
}

impl Mlp {
    fn forward(&self, x: &Array) -> Array {
        self.out.forward(&relu(&self.hidden.forward(x)))
    }
}

random::seed(0);
let mut model = Mlp { hidden: Linear::new(2, 8), out: Linear::new(8, 2) };
assert_eq!(model.paths(), ["hidden.weight", "hidden.bias", "out.weight", "out.bias"]);

let (x, y) = (array![[1.0, -1.0], [0.5, 2.0]], array![0u32, 1]);
let loss = |m: &Mlp| losses::cross_entropy(&m.forward(&x), &y, losses::Reduction::Mean);
let mut opt = Sgd::new(0.5);
let (before, _) = value_and_grad(loss)(&model)?;
for _ in 0..20 {
    let (_, grads) = value_and_grad(loss)(&model)?;   // grads: Mlp
    opt.update(&mut model, &grads)?;
}
let (after, _) = value_and_grad(loss)(&model)?;
assert!(after.item::<f32>()? < before.item::<f32>()?);

The derive is syntactic: a field typed Array (or Option<Array>, Vec<Array>) needs a role (Param for a parameter, #[tree(buffer)], #[tree(skip)] for a computed constant); a type alias of Array is not recognised. Trait-object fields are not trees: use an enum of the concrete layer types.

Stability: no promise when this crate is depended on directly; depend on nuvec, whose re-exports are semver-tracked.

Status: pre-alpha.

Re-exports

pub use layers::Dropout;

pub use layers::Embedding;

pub use layers::LayerNorm;

pub use layers::Linear;

pub use layers::LoraBase;

pub use layers::LoraLinear;

pub use layers::QuantizedEmbedding;

pub use layers::QuantizedLinear;

pub use layers::RmsNorm;

pub use layers::Rope;

Modules

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.

losses

Loss functions, as in MLX’s nn.losses.

optim

Optimizers, learning-rate schedules and gradient clipping, with MLX’s formulas (mlx.optimizers).

sharded

Tensor-parallel linear layers over a distributed group, as mlx.nn’s AllToShardedLinear and ShardedToAllLinear (and their quantized forms).

Structs

Llama3

RopeScaling::Llama3’s parameters (mlx-lm’s defaults).

LongRope

RopeScaling::LongRope’s parameters (mlx-lm’s defaults, long factors).

ModuleDisplay

A model’s nested modules with their parameter counts (crate::Module::display):

Param

A module parameter: an array and a frozen flag (trainable unless frozen).

Summary

A table of a model’s leaves (path, shape, dtype, trainable, bytes) with totals, printed as aligned columns (crate::Module::summary).

SummaryRow

One leaf of a Summary.

Yarn

RopeScaling::Yarn’s parameters (mlx-lm’s defaults; Hugging Face’s truncate is ignored, as by mlx-lm).

Enums

RopeScaling

How a rotary embedding’s frequencies (or positions) are scaled, as mlx-lm’s RoPE classes. Frequencies are computed in float32 with mlx-lm’s formulas.

Constants

VERSION

The version of this crate.

Traits

Module

Model operations, as MLX’s nn.Module, for every Tree (a blanket impl): a model is a struct of layers with #[derive(Module)], and its leaves’ paths are MLX’s names.

Tree

An owned tree of arrays: rebuilt with new leaves of the same structure, mutated in place.

WeightSource

Where crate::Module::load_weights reads named weights from: a safetensors file (a path), or named arrays (Vec<(String, Array)>, a HashMap or BTreeMap).

Functions

geglu

gelu_approx(gate) · x in one kernel, as mlx-lm’s compiled geglu (Gemma 4).

gelu_approx

GELU with the tanh approximation, as MLX’s nn.gelu_approx: 0.5 x (1 + tanh(√(2/π) (x + 0.044715 x³))).

gelu_fast_approx

x · sigmoid(1.702 x), as MLX’s nn.gelu_fast_approx.

logit_softcap

tanh(x / cap) · cap in one kernel, as mlx-lm’s compiled logit_softcap (Gemma 4’s final logits: cap 30). A cap that is not positive and finite gives a failed array.

relu

max(x, 0), as MLX’s nn.relu.

silu

x · sigmoid(x), as MLX’s nn.silu.

swiglu

silu(gate) · x in one kernel, as mlx-lm’s compiled swiglu.

swiglu_oai

gpt-oss’s clamped SwiGLU in one kernel, as mlx-lm’s compiled swiglu_oai: (clip(x, −limit, limit) + 1) · g · sigmoid(alpha · g) with g = min(gate, limit) (gpt-oss: alpha 1.702, limit 7).

Derive Macros

Module

Derives a model’s Tree impl, with a check against bare Array fields (see the crate documentation). Every tree is a Module (a blanket impl). Derives Tree for a model (it emits no impl Module: Module is implemented for every tree), with one more check: a field typed Array (directly or in Option, Vec or Box) without buffer or skip is an error suggesting Param. The check is syntactic: a type alias of Array is not caught.

Tree

Derives Leaves and Tree for a struct (named, tuple or unit) or an enum: the type’s array leaves get paths from its field names, as MLX’s tree_flatten.