nuvec-nn 0.6.1

//! 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' `forward`s, 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)))
//!     }
//! }
//!
//! # fn main() -> Result<()> {
//! 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>()?);
//! # Ok(())
//! # }
//! ```
//!
//! 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.

#![forbid(unsafe_code)]

use std::any::Any;
use std::path::Path;

use nuvec::tree::{LeafKind, LoadMode, Visit, VisitMut, Visitor, VisitorMut};
use nuvec::{Array, Category, Error, Result};

mod display;
pub mod layers;
pub mod losses;
pub mod optim;
mod rope_scaling;
pub mod sharded;
mod weights;

pub use display::{ModuleDisplay, Summary, SummaryRow};
pub use layers::{
    Dropout, Embedding, LayerNorm, Linear, LoraBase, LoraLinear, QuantizedEmbedding,
    QuantizedLinear, RmsNorm, Rope,
};
/// 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).
pub use nuvec::macros::Module;
pub use nuvec::{Param, Tree};
pub use rope_scaling::{Llama3, LongRope, RopeScaling, Yarn};
pub use weights::WeightSource;

/// The Rust snippets of the user guide's neural-network and distributed pages, run as
/// doctests.
#[cfg(doctest)]
mod site_doctests {
    #[doc = include_str!("../../../site/content/docs/neural-networks.md")]
    pub struct NeuralNetworks;
    #[doc = include_str!("../../../site/content/docs/distributed.md")]
    pub struct Distributed;
}

/// The version of this crate.
pub const VERSION: &str = env!("CARGO_PKG_VERSION");

/// 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.
///
/// Leaves are parameters ([`Param`], trainable unless frozen), buffers (never trained,
/// such as quantized weights) or plain arrays.
pub trait Module: Tree {
    /// Every leaf (parameters, buffers and plain arrays) with its path, in order.
    fn parameters(&self) -> Vec<(String, Array)> {
        self.flatten()
    }

    /// The trainable leaves with their paths, in order.
    fn trainable_parameters(&self) -> Vec<(String, Array)> {
        let mut out = Vec::new();
        visit_leaves(self, |p, a, _, trainable| {
            if trainable {
                out.push((p.to_owned(), a.clone()));
            }
        });
        out
    }

    /// The number of elements over every leaf (MLX's `tree_flatten` total).
    fn num_params(&self) -> usize {
        let mut n = 0;
        visit_leaves(self, |_, a, _, _| n += a.size());
        n
    }

    /// The number of elements over the trainable leaves.
    fn num_trainable(&self) -> usize {
        let mut n = 0;
        visit_leaves(self, |_, a, _, trainable| {
            if trainable {
                n += a.size();
            }
        });
        n
    }

    /// A table of the leaves (path, shape, dtype, trainable, bytes) with totals.
    fn summary(&self) -> Summary {
        Summary::of(self)
    }

    /// The nested modules with their parameter counts, for printing:
    /// `println!("{}", model.display())`.
    fn display(&self) -> ModuleDisplay<'_, Self> {
        ModuleDisplay::new(self)
    }

    /// Freezes the parameters whose path matches `filter`, returning how many changed.
    /// A matching plain `Array` leaf cannot be frozen (only [`Param`]s can): that is an
    /// error naming the paths, and nothing changes. Buffers are never trainable anyway.
    fn freeze(&mut self, filter: impl Fn(&str) -> bool) -> Result<usize> {
        set_frozen(self, &filter, true)
    }

    /// Unfreezes the parameters whose path matches `filter`, returning how many changed.
    /// Buffers (quantized weights) and integer parameters stay non-trainable.
    fn unfreeze(&mut self, filter: impl Fn(&str) -> bool) -> Result<usize> {
        set_frozen(self, &filter, false)
    }

    /// Sets every mode flag of the tree: training (`true`, the default) or evaluation
    /// (`false`, where dropout is the identity), as MLX's `Module.train(mode)`. (There is
    /// no `eval()`: it would read like `x.eval()`, which evaluates.)
    fn train(&mut self, on: bool) {
        struct Train(bool);
        impl VisitorMut for Train {
            fn leaf(&mut self, _: &str, _: &mut Array, _: LeafKind, _: Option<&mut bool>) {}
            fn training(&mut self, _: &str, flag: &mut bool) {
                *flag = self.0;
            }
        }
        self.visit_mut(&mut VisitMut::new(&mut Train(on)));
    }

    /// Whether the tree is in training mode (its first mode flag; `true` without any).
    fn training(&self) -> bool {
        struct Mode(Option<bool>);
        impl Visitor for Mode {
            fn leaf(&mut self, _: &str, _: &Array, _: LeafKind, _: bool) {}
            fn mode(&mut self, _: &str, training: bool) {
                self.0.get_or_insert(training);
            }
        }
        let mut m = Mode(None);
        self.visit(&mut Visit::new(&mut m));
        m.0.unwrap_or(true)
    }

    /// Replaces leaves by name, strictly (MLX's `load_weights(strict=True)`): every leaf
    /// must be given with its shape and nothing else (loaded dtypes are kept). On any
    /// problem the error lists every missing, unknown and mismatched name, and the model
    /// is unchanged. `src` is a safetensors path or named arrays.
    fn load_weights(&mut self, src: impl WeightSource) -> Result<()> {
        weights::load(self, src, LoadMode::Strict)
    }

    /// Replaces the leaves given by name (MLX's `load_weights(strict=False)`): unknown
    /// names are ignored, missing leaves keep their values, shapes must match.
    fn load_weights_partial(&mut self, src: impl WeightSource) -> Result<()> {
        weights::load(self, src, LoadMode::Partial)
    }

    /// Writes every leaf as safetensors with MLX's names and dtypes (MLX's
    /// `save_weights`), evaluating them. Unloaded placeholders are an error naming them.
    fn save_weights(&self, path: impl AsRef<Path>) -> Result<()> {
        weights::save(self, path.as_ref(), false)
    }

    /// Writes only the trainable leaves (an adapter's weights, say).
    fn save_trainable(&self, path: impl AsRef<Path>) -> Result<()> {
        weights::save(self, path.as_ref(), true)
    }

    /// Calls `f(path, submodule)` on every submodule of type `T` whose path matches
    /// `filter` (`f` may replace it, as when applying LoRA), returning how many. It stops at
    /// the first error from `f`; submodules changed before it keep their change. Matching
    /// submodules are not searched for further matches.
    fn map_modules<T: Tree>(
        &mut self,
        filter: impl Fn(&str) -> bool,
        mut f: impl FnMut(&str, &mut T) -> Result<()>,
    ) -> Result<usize> {
        struct Map<'a, T> {
            filter: &'a dyn Fn(&str) -> bool,
            f: &'a mut dyn FnMut(&str, &mut T) -> Result<()>,
            count: usize,
            error: Option<Error>,
        }
        impl<T: 'static> VisitorMut for Map<'_, T> {
            fn leaf(&mut self, _: &str, _: &mut Array, _: LeafKind, _: Option<&mut bool>) {}
            fn child(&mut self, path: &str, child: &mut dyn Any) -> bool {
                if self.error.is_some() {
                    return false;
                }
                match child.downcast_mut::<T>() {
                    Some(m) if (self.filter)(path) => {
                        match (self.f)(path, m) {
                            Ok(()) => self.count += 1,
                            Err(e) => self.error = Some(e),
                        }
                        false
                    }
                    _ => true,
                }
            }
        }
        let mut map = Map {
            filter: &filter,
            f: &mut f,
            count: 0,
            error: None,
        };
        self.visit_mut(&mut VisitMut::new(&mut map));
        match map.error {
            Some(e) => Err(e),
            None => Ok(map.count),
        }
    }
}

impl<T: Tree> Module for T {}

/// Calls `f(path, leaf, kind, trainable)` for every leaf.
fn visit_leaves<T: Tree>(tree: &T, f: impl FnMut(&str, &Array, LeafKind, bool)) {
    struct V<F>(F);
    impl<F: FnMut(&str, &Array, LeafKind, bool)> Visitor for V<F> {
        fn leaf(&mut self, path: &str, a: &Array, kind: LeafKind, trainable: bool) {
            (self.0)(path, a, kind, trainable);
        }
    }
    tree.visit(&mut Visit::new(&mut V(f)));
}

fn set_frozen<T: Tree>(tree: &mut T, filter: &dyn Fn(&str) -> bool, frozen: bool) -> Result<usize> {
    if frozen {
        let mut plain = Vec::new();
        visit_leaves(tree, |p, _, kind, _| {
            if kind == LeafKind::Array && filter(p) {
                plain.push(p.to_owned());
            }
        });
        if !plain.is_empty() {
            return Err(Error::InvalidArgument(format!(
                "freeze: plain arrays cannot be frozen (make them `Param`s): {}",
                plain.join(", ")
            )));
        }
    }
    struct Set<'a> {
        filter: &'a dyn Fn(&str) -> bool,
        frozen: bool,
        changed: usize,
    }
    impl VisitorMut for Set<'_> {
        fn leaf(&mut self, path: &str, a: &mut Array, kind: LeafKind, flag: Option<&mut bool>) {
            let Some(flag) = flag else { return };
            if kind != LeafKind::Param || !(self.filter)(path) || *flag == self.frozen {
                return;
            }
            // Integer parameters (packed weights, say) are never trained.
            if !self.frozen && !matches!(a.dtype().category(), Category::Float | Category::Complex)
            {
                return;
            }
            *flag = self.frozen;
            self.changed += 1;
        }
    }
    let mut set = Set {
        filter,
        frozen,
        changed: 0,
    };
    tree.visit_mut(&mut VisitMut::new(&mut set));
    Ok(set.changed)
}

/// `max(x, 0)`, as MLX's `nn.relu`.
#[track_caller]
pub fn relu(x: &Array) -> Array {
    x.maximum(0.0)
}

/// `x · sigmoid(x)`, as MLX's `nn.silu`.
#[track_caller]
pub fn silu(x: &Array) -> Array {
    x.multiply(x.sigmoid())
}

/// GELU with the tanh approximation, as MLX's `nn.gelu_approx`:
/// `0.5 x (1 + tanh(√(2/π) (x + 0.044715 x³)))`.
#[track_caller]
pub fn gelu_approx(x: &Array) -> Array {
    // `x³` as MLX's `x**3` (one power, one rounding; not two products).
    let cube = x.power(3.0);
    let inner = x
        .add(cube.multiply(0.044715))
        .multiply((2.0 / std::f64::consts::PI).sqrt());
    x.multiply(0.5).multiply(inner.tanh().add(1.0))
}

/// `x · sigmoid(1.702 x)`, as MLX's `nn.gelu_fast_approx`.
#[track_caller]
pub fn gelu_fast_approx(x: &Array) -> Array {
    x.multiply(x.multiply(1.702).sigmoid())
}

/// A compiled function of a tuple of arrays (traced once per input signature, element-wise
/// chains fused into one kernel), as mlx-lm compiles its activations.
type Compiled<T> = Box<dyn Fn(&T) -> Result<Array> + Send + Sync>;

/// `x` as a 0-d array of `like`'s dtype: a weak scalar (MLX rounds a Python float to the
/// array's dtype), passed to compiled functions as an input so one trace serves every value.
fn weak(value: f32, like: &Array) -> Array {
    Array::from(value).astype(like.dtype())
}

/// `silu(gate) · x` in one kernel, as mlx-lm's compiled `swiglu`.
#[track_caller]
pub fn swiglu(gate: &Array, x: &Array) -> Array {
    static F: std::sync::LazyLock<Compiled<(Array, Array)>> = std::sync::LazyLock::new(|| {
        Box::new(nuvec::compile(|(g, x): &(Array, Array)| {
            silu(g).multiply(x)
        }))
    });
    F(&(gate.clone(), x.clone())).unwrap_or_else(|_| silu(gate).multiply(x))
}

/// 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).
///
/// ```
/// use nuvec::array;
///
/// let y = nuvec_nn::swiglu_oai(&array![10.0f32, -1.0], &array![9.0f32, 0.5], 1.702, 7.0);
/// let v = y.to_vec::<f32>().unwrap();
/// assert!((v[0] - 8.0 * 7.0 / (1.0 + (-1.702f32 * 7.0).exp())).abs() < 1e-4);
/// ```
#[track_caller]
pub fn swiglu_oai(gate: &Array, x: &Array, alpha: f32, limit: f32) -> Array {
    fn graph(g: &Array, x: &Array, alpha: &Array, limit: &Array) -> Array {
        let g = g.minimum(limit);
        let x = x.maximum(limit.negative()).minimum(limit);
        x.add(1.0)
            .multiply(g.multiply(alpha.multiply(&g).sigmoid()))
    }
    static F: std::sync::LazyLock<Compiled<(Array, Array, Array, Array)>> =
        std::sync::LazyLock::new(|| {
            Box::new(nuvec::compile(
                |(g, x, a, l): &(Array, Array, Array, Array)| graph(g, x, a, l),
            ))
        });
    let (a, l) = (weak(alpha, gate), weak(limit, gate));
    F(&(gate.clone(), x.clone(), a.clone(), l.clone())).unwrap_or_else(|_| graph(gate, x, &a, &l))
}

/// `gelu_approx(gate) · x` in one kernel, as mlx-lm's compiled `geglu` (Gemma 4).
#[track_caller]
pub fn geglu(gate: &Array, x: &Array) -> Array {
    static F: std::sync::LazyLock<Compiled<(Array, Array)>> = std::sync::LazyLock::new(|| {
        Box::new(nuvec::compile(|(g, x): &(Array, Array)| {
            gelu_approx(g).multiply(x)
        }))
    });
    F(&(gate.clone(), x.clone())).unwrap_or_else(|_| gelu_approx(gate).multiply(x))
}

/// `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.
///
/// No attention option caps scores; a soft-capped attention is composed from public ops:
///
/// ```
/// use nuvec::{DType, ops, random};
///
/// let (q, k, v) = (random::normal([1, 2, 4, 8], DType::F32, 0.0, 1.0, None),
///                  random::normal([1, 2, 6, 8], DType::F32, 0.0, 1.0, None),
///                  random::normal([1, 2, 6, 8], DType::F32, 0.0, 1.0, None));
/// let (scale, cap) = (8f32.sqrt().recip(), 50.0);
/// let scores = q.matmul(&k.transpose([0, 1, 3, 2])).multiply(scale);
/// let probs = nuvec_nn::logit_softcap(&scores, cap).softmax(-1); // precise for half types
/// let out = probs.matmul(&v);
/// assert_eq!(out.shape(), [1, 2, 4, 8]);
/// ```
#[track_caller]
pub fn logit_softcap(x: &Array, cap: f32) -> Array {
    if !(cap > 0.0 && cap.is_finite()) {
        return Array::failure(
            "logit_softcap",
            nuvec::OpErrorKind::InvalidArgument(format!(
                "the cap must be positive and finite, got {cap}"
            )),
        );
    }
    fn graph(x: &Array, cap: &Array) -> Array {
        x.divide(cap).tanh().multiply(cap)
    }
    static F: std::sync::LazyLock<Compiled<(Array, Array)>> = std::sync::LazyLock::new(|| {
        Box::new(nuvec::compile(|(x, c): &(Array, Array)| graph(x, c)))
    });
    let c = weak(cap, x);
    F(&(x.clone(), c.clone())).unwrap_or_else(|_| graph(x, &c))
}