Crate nuvec_nn 

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

• Module: named parameter trees (layers.0.weight), trainable or frozen, loaded from checkpoints by name (Module::load_weights).
• Layers (layers): Linear, QuantizedLinear, Embedding, QuantizedEmbedding, RmsNorm, LayerNorm, Rope and LoraLinear; activations relu, silu, gelu_approx and gelu_fast_approx.
• losses::cross_entropy, value_and_grad for a model’s loss, and optim: SGD, Adam, AdamW, Lion and Muon with learning-rate schedules and gradient clipping.

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

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

impl Module for Mlp {
    fn visit(&mut self, p: &mut Parameters<'_>) {
        p.module("hidden", &mut self.hidden);
        p.module("out", &mut self.out);
    }
}

random::seed(0);
let mut model = Mlp { hidden: Linear::new(2, 8, true)?, out: Linear::new(8, 2, true)? };
let names: Vec<String> = model.parameters().into_iter().map(|(n, _)| n).collect();
assert_eq!(names, ["hidden.weight", "hidden.bias", "out.weight", "out.bias"]);

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

Status: pre-alpha.

Re-exports

pub use layers::Embedding;

pub use layers::LayerNorm;

pub use layers::Linear;

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

Structs

Frozen

A module whose parameters are all frozen (MLX’s freeze).

Parameters

Receives a module’s parameters during Module::visit, with their paths.

Constants

VERSION

The version of this crate.

Traits

Module

A model: a tree of named parameters, as MLX’s nn.Module.

Functions

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.

relu

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

silu

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

value_and_grad

The loss of model and its gradient with respect to every trainable parameter (in the order of Module::trainable_parameters), as MLX’s nn.value_and_grad.