geographdb_core/algorithms/

mlp.rs

//! Pure-Rust MLP primitives for small geometric training loops.
//!
//! No external BLAS or autograd framework — just row-major matrix multiply,
//! stable softmax, cross-entropy, and a tiny two-layer MLP with manual
//! back-propagation. Intended for CPU demos and experiments that need a
//! trainable geometric head without pulling in a full ML framework.

use crate::algorithms::natural_grad::softmax;
use ndarray::{Array1, Array2, ArrayView2, Axis};

/// Row-major matrix multiplication.
///
/// `a` is `m × k`, `b` is `k × n`, and the returned vector is `m × n`.
pub fn matmul(a: &[f32], a_rows: usize, a_cols: usize, b: &[f32], b_cols: usize) -> Vec<f32> {
    assert_eq!(a.len(), a_rows * a_cols, "matmul: a dimensions mismatch");
    assert_eq!(b.len(), a_cols * b_cols, "matmul: b dimensions mismatch");

    let mut c = vec![0.0f32; a_rows * b_cols];
    matmul_into(a, a_rows, a_cols, b, b_cols, &mut c);
    c
}

/// Row-major matrix multiplication writing into a pre-allocated `c` buffer.
///
/// `c` must have length `a_rows * b_cols`.
pub fn matmul_into(
    a: &[f32],
    a_rows: usize,
    a_cols: usize,
    b: &[f32],
    b_cols: usize,
    c: &mut [f32],
) {
    assert_eq!(
        a.len(),
        a_rows * a_cols,
        "matmul_into: a dimensions mismatch"
    );
    assert_eq!(
        b.len(),
        a_cols * b_cols,
        "matmul_into: b dimensions mismatch"
    );
    assert_eq!(
        c.len(),
        a_rows * b_cols,
        "matmul_into: c dimensions mismatch"
    );

    // matrixmultiply::sgemm gives near-BLAS performance for f32.
    // Strides are row-major: row stride = width, col stride = 1.
    unsafe {
        matrixmultiply::sgemm(
            a_rows,
            a_cols,
            b_cols,
            1.0,
            a.as_ptr(),
            a_cols as isize,
            1,
            b.as_ptr(),
            b_cols as isize,
            1,
            0.0,
            c.as_mut_ptr(),
            b_cols as isize,
            1,
        );
    }
}

/// Cross-entropy loss from a probability distribution and an integer target class.
///
/// Returns `-ln(probs[target])` with a small floor to avoid `log(0)`.
pub fn cross_entropy_loss(probs: &[f32], target: usize) -> f32 {
    assert!(!probs.is_empty(), "cross_entropy_loss: empty distribution");
    assert!(
        target < probs.len(),
        "cross_entropy_loss: target out of bounds"
    );
    -probs[target].max(1e-30).ln()
}

/// Stable cross-entropy directly from logits and an integer target class.
///
/// Computes `log(sum(exp(logits))) - logits[target]` in a numerically stable way
/// so large-magnitude logits do not overflow.
pub fn cross_entropy_from_logits(logits: &[f32], target: usize) -> f32 {
    assert!(
        !logits.is_empty(),
        "cross_entropy_from_logits: empty logits"
    );
    assert!(
        target < logits.len(),
        "cross_entropy_from_logits: target out of bounds"
    );
    let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let log_sum_exp = logits.iter().map(|&l| (l - max).exp()).sum::<f32>().ln() + max;
    -logits[target] + log_sum_exp
}

/// Gradient of the cross-entropy loss with respect to logits.
///
/// Returns `softmax(logits) - one_hot(target)`.
pub fn cross_entropy_logits_grad(logits: &[f32], target: usize) -> Vec<f32> {
    let mut probs = softmax(logits);
    probs[target] -= 1.0;
    probs
}

/// Build a one-hot column vector of length `classes` with a 1 at `class`.
pub fn one_hot(class: usize, classes: usize) -> Vec<f32> {
    assert!(class < classes, "one_hot: class out of bounds");
    let mut v = vec![0.0f32; classes];
    v[class] = 1.0;
    v
}

/// Rectified linear unit.
pub fn relu(x: f32) -> f32 {
    x.max(0.0)
}

/// Derivative of ReLU.
pub fn relu_grad(x: f32) -> f32 {
    if x > 0.0 {
        1.0
    } else {
        0.0
    }
}

/// Tiny deterministic xorshift RNG used for weight initialisation.
struct XorShift32 {
    state: u32,
}

impl XorShift32 {
    fn new(seed: u32) -> Self {
        Self {
            state: seed.wrapping_add(0x9e37_79b9),
        }
    }

    fn next(&mut self) -> u32 {
        let mut x = self.state;
        x ^= x << 13;
        x ^= x >> 17;
        x ^= x << 5;
        self.state = x;
        x
    }

    /// Uniform random float in `(-scale, scale)`.
    fn uniform_f32(&mut self, scale: f32) -> f32 {
        let u = self.next();
        let normalised = (u as f32 / u32::MAX as f32) * 2.0 - 1.0;
        normalised * scale
    }
}

/// Forward-pass cache for a single input.
pub struct MlpForward {
    /// Pre-activation of the hidden layer (`z1`).
    pub z1: Vec<f32>,
    /// Activated hidden layer (`h`).
    pub hidden: Vec<f32>,
    /// Output logits before softmax.
    pub logits: Vec<f32>,
}

/// Forward-pass cache for a batch of inputs.
pub struct MlpForwardBatch {
    /// Pre-activation of the hidden layer, shape `[batch, hidden_dim]`.
    pub z1: Array2<f32>,
    /// Activated hidden layer, shape `[batch, hidden_dim]`.
    pub hidden: Array2<f32>,
    /// Output logits before softmax, shape `[batch, output_dim]`.
    pub logits: Array2<f32>,
}

/// Gradients for a single training example.
#[derive(Debug, Clone)]
pub struct MlpGradients {
    /// `input_dim × hidden_dim`
    pub dw1: Vec<f32>,
    /// `hidden_dim`
    pub db1: Vec<f32>,
    /// `hidden_dim × output_dim`
    pub dw2: Vec<f32>,
    /// `output_dim`
    pub db2: Vec<f32>,
}

/// A small two-layer MLP classifier: input → linear → ReLU → linear → logits.
///
/// Weights are stored row-major:
/// - `w1` has shape `[input_dim, hidden_dim]`
/// - `w2` has shape `[hidden_dim, output_dim]`
#[derive(Debug, Clone)]
pub struct MlpClassifier {
    input_dim: usize,
    hidden_dim: usize,
    output_dim: usize,
    /// Input-to-hidden weights, shape `[input_dim, hidden_dim]`.
    pub w1: Vec<f32>,
    /// Hidden bias, length `hidden_dim`.
    pub b1: Vec<f32>,
    /// Hidden-to-output weights, shape `[hidden_dim, output_dim]`.
    pub w2: Vec<f32>,
    /// Output bias, length `output_dim`.
    pub b2: Vec<f32>,
}

impl MlpClassifier {
    /// Create a fresh classifier with Xavier-like random weights.
    pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize, seed: u32) -> Self {
        let mut rng = XorShift32::new(seed);

        let scale1 = (2.0 / input_dim as f32).sqrt() * 0.1;
        let mut w1 = vec![0.0f32; input_dim * hidden_dim];
        for v in w1.iter_mut() {
            *v = rng.uniform_f32(scale1);
        }
        let mut b1 = vec![0.0f32; hidden_dim];
        for v in b1.iter_mut() {
            *v = rng.uniform_f32(0.01);
        }

        let scale2 = (2.0 / hidden_dim as f32).sqrt() * 0.1;
        let mut w2 = vec![0.0f32; hidden_dim * output_dim];
        for v in w2.iter_mut() {
            *v = rng.uniform_f32(scale2);
        }
        let mut b2 = vec![0.0f32; output_dim];
        for v in b2.iter_mut() {
            *v = rng.uniform_f32(0.01);
        }

        Self {
            input_dim,
            hidden_dim,
            output_dim,
            w1,
            b1,
            w2,
            b2,
        }
    }

    /// Forward pass for a single input vector.
    pub fn forward(&self, x: &[f32]) -> MlpForward {
        assert_eq!(x.len(), self.input_dim, "forward: input dimension mismatch");

        let z1 = matmul(x, 1, self.input_dim, &self.w1, self.hidden_dim)
            .iter()
            .zip(self.b1.iter())
            .map(|(&z, &b)| z + b)
            .collect::<Vec<_>>();
        let hidden = z1.iter().map(|&z| relu(z)).collect::<Vec<_>>();
        let logits = matmul(&hidden, 1, self.hidden_dim, &self.w2, self.output_dim)
            .iter()
            .zip(self.b2.iter())
            .map(|(&z, &b)| z + b)
            .collect::<Vec<_>>();

        MlpForward { z1, hidden, logits }
    }

    /// Hidden state (post-ReLU) for a single input vector, without the
    /// hidden-to-output projection.  Useful for FFN-head diagnostics and
    /// confidence-gated routing where logits may not be needed.
    pub fn forward_hidden(&self, x: &[f32]) -> Vec<f32> {
        assert_eq!(
            x.len(),
            self.input_dim,
            "forward_hidden: input dimension mismatch"
        );

        matmul(x, 1, self.input_dim, &self.w1, self.hidden_dim)
            .iter()
            .zip(self.b1.iter())
            .map(|(&z, &b)| relu(z + b))
            .collect()
    }

    /// Predicted class (argmax over logits).
    pub fn predict(&self, x: &[f32]) -> usize {
        let fwd = self.forward(x);
        fwd.logits
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(i, _)| i)
            .unwrap_or(0)
    }

    /// Class probabilities for a single input vector.
    pub fn probabilities(&self, x: &[f32]) -> Vec<f32> {
        let fwd = self.forward(x);
        softmax(&fwd.logits)
    }

    /// Cross-entropy loss for a single input/target pair.
    pub fn loss(&self, x: &[f32], target: usize) -> f32 {
        let fwd = self.forward(x);
        cross_entropy_from_logits(&fwd.logits, target)
    }

    /// Back-propagation for a single input/target pair.
    pub fn backward(&self, x: &[f32], target: usize) -> MlpGradients {
        assert_eq!(
            x.len(),
            self.input_dim,
            "backward: input dimension mismatch"
        );
        assert!(
            target < self.output_dim,
            "backward: target class out of bounds"
        );

        let fwd = self.forward(x);
        let dlogits = cross_entropy_logits_grad(&fwd.logits, target);

        // Gradients for W2 and b2.
        let mut dw2 = vec![0.0f32; self.hidden_dim * self.output_dim];
        for (i, &h_i) in fwd.hidden.iter().enumerate() {
            for (k, &dlogit_k) in dlogits.iter().enumerate() {
                dw2[i * self.output_dim + k] = h_i * dlogit_k;
            }
        }
        let db2 = dlogits.clone();

        // Back-propagate through W2.
        let mut dh = vec![0.0f32; self.hidden_dim];
        for (i, dh_i) in dh.iter_mut().enumerate() {
            let mut acc = 0.0f32;
            for (k, &dlogit_k) in dlogits.iter().enumerate() {
                acc += dlogit_k * self.w2[i * self.output_dim + k];
            }
            *dh_i = acc;
        }

        // ReLU derivative.
        let dz1: Vec<f32> = dh
            .iter()
            .zip(fwd.z1.iter())
            .map(|(&dh_i, &z1_i)| dh_i * relu_grad(z1_i))
            .collect();

        // Gradients for W1 and b1.
        let mut dw1 = vec![0.0f32; self.input_dim * self.hidden_dim];
        for j in 0..self.input_dim {
            for i in 0..self.hidden_dim {
                dw1[j * self.hidden_dim + i] = x[j] * dz1[i];
            }
        }
        let db1 = dz1;

        MlpGradients { dw1, db1, dw2, db2 }
    }

    /// Apply a single SGD step using the supplied gradients.
    pub fn apply_sgd(&mut self, grad: &MlpGradients, lr: f32) {
        for i in 0..self.w1.len() {
            self.w1[i] -= lr * grad.dw1[i];
        }
        for i in 0..self.b1.len() {
            self.b1[i] -= lr * grad.db1[i];
        }
        for i in 0..self.w2.len() {
            self.w2[i] -= lr * grad.dw2[i];
        }
        for i in 0..self.b2.len() {
            self.b2[i] -= lr * grad.db2[i];
        }
    }

    /// Forward pass for a batch of inputs.
    ///
    /// `x` is a row-major flat buffer of shape `[batch_size, input_dim]`.
    pub fn forward_batch(&self, x: &[f32], batch_size: usize) -> MlpForwardBatch {
        assert_eq!(
            x.len(),
            batch_size * self.input_dim,
            "forward_batch: input dimension mismatch"
        );

        let x_arr = Array2::from_shape_vec((batch_size, self.input_dim), x.to_vec())
            .expect("forward_batch: invalid input shape");
        let w1_view = ArrayView2::from_shape((self.input_dim, self.hidden_dim), &self.w1).unwrap();
        let b1_view = ArrayView2::from_shape((1, self.hidden_dim), &self.b1).unwrap();
        let z1 = x_arr.dot(&w1_view) + &b1_view;

        let hidden = z1.mapv(|v| relu(v));

        let w2_view = ArrayView2::from_shape((self.hidden_dim, self.output_dim), &self.w2).unwrap();
        let b2_view = ArrayView2::from_shape((1, self.output_dim), &self.b2).unwrap();
        let logits = hidden.dot(&w2_view) + &b2_view;

        MlpForwardBatch { z1, hidden, logits }
    }

    /// Cross-entropy loss averaged over a batch.
    pub fn loss_batch(&self, x: &[f32], targets: &[usize], batch_size: usize) -> f32 {
        assert_eq!(
            targets.len(),
            batch_size,
            "loss_batch: target count mismatch"
        );
        let fwd = self.forward_batch(x, batch_size);
        let mut total = 0.0f32;
        for (logits_row, &target) in fwd.logits.axis_iter(Axis(0)).zip(targets.iter()) {
            total += cross_entropy_from_logits(logits_row.as_slice().unwrap(), target);
        }
        total / batch_size as f32
    }

    /// Back-propagation for a batch of input/target pairs.
    ///
    /// Returns gradients **averaged over the batch**, the input gradient
    /// `dx` as a row-major `[batch_size, input_dim]` buffer, and the average
    /// cross-entropy loss. A single SGD/Adam step with the returned gradients
    /// is equivalent to one step per mini-batch (not one step per example).
    pub fn backward_batch(
        &self,
        x: &[f32],
        targets: &[usize],
        batch_size: usize,
    ) -> (MlpGradients, Vec<f32>, f32) {
        assert_eq!(
            x.len(),
            batch_size * self.input_dim,
            "backward_batch: input dimension mismatch"
        );
        assert_eq!(
            targets.len(),
            batch_size,
            "backward_batch: target count mismatch"
        );
        for &target in targets {
            assert!(
                target < self.output_dim,
                "backward_batch: target class out of bounds"
            );
        }

        let fwd = self.forward_batch(x, batch_size);
        let batch_f = batch_size as f32;

        // dlogits = softmax(logits) - one_hot(target), averaged implicitly later.
        let mut dlogits = Array2::zeros((batch_size, self.output_dim));
        let mut total_loss = 0.0f32;
        for ((mut row, logits_row), &target) in dlogits
            .axis_iter_mut(Axis(0))
            .zip(fwd.logits.axis_iter(Axis(0)))
            .zip(targets.iter())
        {
            let probs = softmax(logits_row.as_slice().unwrap());
            total_loss += -probs[target].max(1e-30).ln();
            let mut shifted = Array1::from_vec(probs);
            shifted[target] -= 1.0;
            row.assign(&shifted);
        }
        let avg_loss = total_loss / batch_f;

        // Gradients for W2 and b2 (averaged over the batch).
        let dw2 = fwd.hidden.t().dot(&dlogits) / batch_f;
        let db2 = dlogits.mean_axis(Axis(0)).unwrap();

        // Back-propagate through W2.
        let w2_view = ArrayView2::from_shape((self.hidden_dim, self.output_dim), &self.w2).unwrap();
        let dh = dlogits.dot(&w2_view.t());

        // ReLU derivative.
        let dz1 = dh * fwd.z1.mapv(|v| relu_grad(v));

        // Input gradient (not averaged; per-example so callers can update
        // external embedding tables with the same batch semantics as the MLP).
        let w1_view = ArrayView2::from_shape((self.input_dim, self.hidden_dim), &self.w1).unwrap();
        let dx = dz1.dot(&w1_view.t());

        // Gradients for W1 and b1 (averaged over the batch).
        let x_arr = Array2::from_shape_vec((batch_size, self.input_dim), x.to_vec())
            .expect("backward_batch: invalid input shape");
        let dw1 = x_arr.t().dot(&dz1) / batch_f;
        let db1 = dz1.mean_axis(Axis(0)).unwrap();

        (
            MlpGradients {
                dw1: dw1.into_raw_vec(),
                db1: db1.into_raw_vec(),
                dw2: dw2.into_raw_vec(),
                db2: db2.into_raw_vec(),
            },
            dx.into_raw_vec(),
            avg_loss,
        )
    }

    /// Flatten all learnable parameters into one vector, ordered
    /// `[w1, b1, w2, b2]`.
    pub fn flatten_params(&self) -> Vec<f32> {
        let mut params =
            Vec::with_capacity(self.w1.len() + self.b1.len() + self.w2.len() + self.b2.len());
        params.extend_from_slice(&self.w1);
        params.extend_from_slice(&self.b1);
        params.extend_from_slice(&self.w2);
        params.extend_from_slice(&self.b2);
        params
    }

    /// Flatten the gradients for a single example in the same order as
    /// [`Self::flatten_params`].
    pub fn flatten_grad(&self, grad: &MlpGradients) -> Vec<f32> {
        let mut flat =
            Vec::with_capacity(grad.dw1.len() + grad.db1.len() + grad.dw2.len() + grad.db2.len());
        flat.extend_from_slice(&grad.dw1);
        flat.extend_from_slice(&grad.db1);
        flat.extend_from_slice(&grad.dw2);
        flat.extend_from_slice(&grad.db2);
        flat
    }

    /// Restore parameters from a flat vector produced by [`Self::flatten_params`].
    pub fn load_flat_params(&mut self, params: &[f32]) {
        let w1_len = self.w1.len();
        let b1_len = self.b1.len();
        let w2_len = self.w2.len();
        let b2_len = self.b2.len();
        let expected = w1_len + b1_len + w2_len + b2_len;
        assert_eq!(params.len(), expected, "load_flat_params: size mismatch");

        let mut off = 0;
        self.w1.copy_from_slice(&params[off..off + w1_len]);
        off += w1_len;
        self.b1.copy_from_slice(&params[off..off + b1_len]);
        off += b1_len;
        self.w2.copy_from_slice(&params[off..off + w2_len]);
        off += w2_len;
        self.b2.copy_from_slice(&params[off..off + b2_len]);
    }
}

/// Gradients for an embedding + MLP classifier.
#[derive(Debug, Clone)]
pub struct EmbeddingMlpGradients {
    /// `vocab_size × embed_dim`
    pub dembedding: Vec<f32>,
    /// MLP weight/bias gradients.
    pub mlp: MlpGradients,
}

/// A two-layer MLP classifier with a learnable token-embedding lookup.
///
/// This avoids materialising one-hot vectors for large vocabularies.  The
/// input is built from `n_context_tokens` token IDs whose embeddings are
/// concatenated, followed by an optional `continuous_dim` vector of continuous
/// features (e.g. geometric positions).
#[derive(Debug, Clone)]
pub struct EmbeddingMlpClassifier {
    vocab_size: usize,
    embed_dim: usize,
    n_context_tokens: usize,
    continuous_dim: usize,
    /// Optional RoPE theta.  When `Some(theta)`, context-token embeddings are
    /// rotated by their sequence distance using standard RoPE frequencies.
    rope_theta: Option<f32>,
    /// Token embedding table, row-major `[vocab_size, embed_dim]`.
    pub embedding: Vec<f32>,
    /// Underlying MLP head.  Its input dimension is
    /// `n_context_tokens * embed_dim + continuous_dim`.
    pub mlp: MlpClassifier,
}

/// Apply standard RoPE rotation to a single embedding vector in place.
///
/// `position` is the 1-based sequence distance from the target (e.g. 1 for the
/// previous token, 2 for the token before that).  `embed_dim` must be even.
fn apply_rope_in_place(vec: &mut [f32], embed_dim: usize, position: usize, theta: f32) {
    assert_eq!(
        vec.len(),
        embed_dim,
        "apply_rope_in_place: dimension mismatch"
    );
    assert!(
        embed_dim % 2 == 0,
        "apply_rope_in_place: embed_dim must be even"
    );

    let ln_theta = theta.ln();
    let pos_f = position as f32;
    let embed_dim_f = embed_dim as f32;

    for pair in 0..embed_dim / 2 {
        let d0 = pair * 2;
        let d1 = d0 + 1;
        let freq = (-2.0f32 * pair as f32 * ln_theta / embed_dim_f).exp();
        let angle = pos_f * freq;
        let (cos_a, sin_a) = (angle.cos(), angle.sin());
        let v0 = vec[d0];
        let v1 = vec[d1];
        vec[d0] = v0 * cos_a - v1 * sin_a;
        vec[d1] = v0 * sin_a + v1 * cos_a;
    }
}

/// Apply the inverse RoPE rotation to a single embedding vector in place.
fn apply_rope_inv_in_place(vec: &mut [f32], embed_dim: usize, position: usize, theta: f32) {
    assert_eq!(
        vec.len(),
        embed_dim,
        "apply_rope_inv_in_place: dimension mismatch"
    );
    assert!(
        embed_dim % 2 == 0,
        "apply_rope_inv_in_place: embed_dim must be even"
    );

    let ln_theta = theta.ln();
    let pos_f = position as f32;
    let embed_dim_f = embed_dim as f32;

    for pair in 0..embed_dim / 2 {
        let d0 = pair * 2;
        let d1 = d0 + 1;
        let freq = (-2.0f32 * pair as f32 * ln_theta / embed_dim_f).exp();
        let angle = pos_f * freq;
        let (cos_a, sin_a) = (angle.cos(), angle.sin());
        let v0 = vec[d0];
        let v1 = vec[d1];
        vec[d0] = v0 * cos_a + v1 * sin_a;
        vec[d1] = -v0 * sin_a + v1 * cos_a;
    }
}

impl EmbeddingMlpClassifier {
    /// Create a fresh classifier with Xavier-like random weights.
    pub fn new(
        vocab_size: usize,
        embed_dim: usize,
        n_context_tokens: usize,
        continuous_dim: usize,
        hidden_dim: usize,
        output_dim: usize,
        seed: u32,
        rope_theta: Option<f32>,
    ) -> Self {
        if let Some(theta) = rope_theta {
            assert!(
                theta > 0.0,
                "EmbeddingMlpClassifier: rope_theta must be positive"
            );
            assert!(
                embed_dim % 2 == 0,
                "EmbeddingMlpClassifier: RoPE requires even embed_dim"
            );
        }

        let input_dim = n_context_tokens * embed_dim + continuous_dim;
        let mut rng = XorShift32::new(seed);

        let scale = (2.0 / embed_dim as f32).sqrt() * 0.1;
        let mut embedding = vec![0.0f32; vocab_size * embed_dim];
        for v in embedding.iter_mut() {
            *v = rng.uniform_f32(scale);
        }

        let mlp = MlpClassifier::new(input_dim, hidden_dim, output_dim, seed.wrapping_add(1));

        Self {
            vocab_size,
            embed_dim,
            n_context_tokens,
            continuous_dim,
            rope_theta,
            embedding,
            mlp,
        }
    }

    /// Total flat input dimension after embedding lookup and concatenation.
    pub fn input_dim(&self) -> usize {
        self.n_context_tokens * self.embed_dim + self.continuous_dim
    }

    fn fill_input(&self, token_ids: &[u32], continuous: Option<&[f32]>, input: &mut [f32]) {
        assert_eq!(
            token_ids.len(),
            self.n_context_tokens,
            "fill_input: token count mismatch"
        );
        if let Some(cont) = continuous {
            assert_eq!(
                cont.len(),
                self.continuous_dim,
                "fill_input: continuous dimension mismatch"
            );
        }

        let mut off = 0;
        for (t, &tid) in token_ids.iter().enumerate() {
            let idx = (tid as usize).min(self.vocab_size - 1);
            let src = &self.embedding[idx * self.embed_dim..(idx + 1) * self.embed_dim];
            input[off..off + self.embed_dim].copy_from_slice(src);

            // RoPE: tokens are ordered [oldest, ..., newest]; newest is at
            // position 1 (distance 1 from target).
            if let Some(theta) = self.rope_theta {
                let position = self.n_context_tokens - t;
                apply_rope_in_place(
                    &mut input[off..off + self.embed_dim],
                    self.embed_dim,
                    position,
                    theta,
                );
            }

            off += self.embed_dim;
        }
        if let Some(cont) = continuous {
            input[off..off + self.continuous_dim].copy_from_slice(cont);
        }
    }

    fn fill_batch_input(
        &self,
        token_ids: &[u32],
        continuous: Option<&[f32]>,
        batch_size: usize,
        input: &mut [f32],
    ) {
        let input_dim = self.input_dim();
        assert_eq!(
            token_ids.len(),
            batch_size * self.n_context_tokens,
            "fill_batch_input: token count mismatch"
        );
        assert_eq!(
            input.len(),
            batch_size * input_dim,
            "fill_batch_input: input buffer size mismatch"
        );

        for b in 0..batch_size {
            let tok_start = b * self.n_context_tokens;
            let cont = continuous.map(|c| {
                let start = b * self.continuous_dim;
                &c[start..start + self.continuous_dim]
            });
            self.fill_input(
                &token_ids[tok_start..tok_start + self.n_context_tokens],
                cont,
                &mut input[b * input_dim..(b + 1) * input_dim],
            );
        }
    }

    /// Forward pass for a single example.
    pub fn forward(&self, token_ids: &[u32], continuous: Option<&[f32]>) -> MlpForward {
        let mut input = vec![0.0f32; self.input_dim()];
        self.fill_input(token_ids, continuous, &mut input);
        self.mlp.forward(&input)
    }

    /// Predicted class for a single example.
    pub fn predict(&self, token_ids: &[u32], continuous: Option<&[f32]>) -> usize {
        let fwd = self.forward(token_ids, continuous);
        fwd.logits
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(i, _)| i)
            .unwrap_or(0)
    }

    /// Cross-entropy loss for a single example.
    pub fn loss(&self, token_ids: &[u32], continuous: Option<&[f32]>, target: usize) -> f32 {
        let fwd = self.forward(token_ids, continuous);
        cross_entropy_from_logits(&fwd.logits, target)
    }

    /// Forward pass for a batch.
    pub fn forward_batch(
        &self,
        token_ids: &[u32],
        continuous: Option<&[f32]>,
        batch_size: usize,
    ) -> MlpForwardBatch {
        let mut input = vec![0.0f32; batch_size * self.input_dim()];
        self.fill_batch_input(token_ids, continuous, batch_size, &mut input);
        self.mlp.forward_batch(&input, batch_size)
    }

    /// Back-propagation for a batch.
    ///
    /// Returns gradients averaged over the batch and the average cross-entropy
    /// loss.  Embedding gradients are accumulated and averaged; the returned
    /// `dembedding` has the same shape as [`Self::embedding`].
    pub fn backward_batch(
        &self,
        token_ids: &[u32],
        continuous: Option<&[f32]>,
        targets: &[usize],
        batch_size: usize,
    ) -> (EmbeddingMlpGradients, f32) {
        assert_eq!(
            token_ids.len(),
            batch_size * self.n_context_tokens,
            "backward_batch: token count mismatch"
        );
        assert_eq!(
            targets.len(),
            batch_size,
            "backward_batch: target count mismatch"
        );
        if let Some(cont) = continuous {
            assert_eq!(
                cont.len(),
                batch_size * self.continuous_dim,
                "backward_batch: continuous size mismatch"
            );
        }

        let mut input = vec![0.0f32; batch_size * self.input_dim()];
        self.fill_batch_input(token_ids, continuous, batch_size, &mut input);

        let (mlp_grad, dx, loss) = self.mlp.backward_batch(&input, targets, batch_size);

        // Accumulate embedding gradients from per-example dx.
        // If RoPE was applied in the forward pass, apply its inverse to dx
        // before accumulating into the static embedding table.
        let mut dembedding = vec![0.0f32; self.embedding.len()];
        let input_dim = self.input_dim();
        for b in 0..batch_size {
            let dx_row = &dx[b * input_dim..(b + 1) * input_dim];
            let tok_start = b * self.n_context_tokens;
            for (t, &tid) in token_ids[tok_start..tok_start + self.n_context_tokens]
                .iter()
                .enumerate()
            {
                let idx = (tid as usize).min(self.vocab_size - 1);
                let mut dx_emb = dx_row[t * self.embed_dim..(t + 1) * self.embed_dim].to_vec();
                if let Some(theta) = self.rope_theta {
                    let position = self.n_context_tokens - t;
                    apply_rope_inv_in_place(&mut dx_emb, self.embed_dim, position, theta);
                }
                let dst = &mut dembedding[idx * self.embed_dim..(idx + 1) * self.embed_dim];
                for (d, g) in dst.iter_mut().zip(dx_emb.iter()) {
                    *d += *g;
                }
            }
        }

        let inv_b = 1.0 / batch_size as f32;
        for v in dembedding.iter_mut() {
            *v *= inv_b;
        }

        (
            EmbeddingMlpGradients {
                dembedding,
                mlp: mlp_grad,
            },
            loss,
        )
    }

    /// Apply a single SGD step using the supplied gradients.
    pub fn apply_sgd(&mut self, grad: &EmbeddingMlpGradients, lr: f32) {
        for i in 0..self.embedding.len() {
            self.embedding[i] -= lr * grad.dembedding[i];
        }
        self.mlp.apply_sgd(&grad.mlp, lr);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn matmul_identity() {
        let a = vec![1.0f32, 2.0, 3.0, 4.0];
        // 2x2 identity
        let i = vec![1.0f32, 0.0, 0.0, 1.0];
        let c = matmul(&a, 2, 2, &i, 2);
        assert_eq!(c, a);
    }

    #[test]
    fn matmul_small() {
        // [1 2]   [5 6]   [19 22]
        // [3 4] * [7 8] = [43 50]
        let a = vec![1.0f32, 2.0, 3.0, 4.0];
        let b = vec![5.0f32, 6.0, 7.0, 8.0];
        let c = matmul(&a, 2, 2, &b, 2);
        assert_eq!(c, vec![19.0, 22.0, 43.0, 50.0]);
    }

    #[test]
    fn cross_entropy_decreases_with_target_probability() {
        let p = vec![0.1f32, 0.8, 0.1];
        assert!(cross_entropy_loss(&p, 1) < cross_entropy_loss(&p, 0));
    }

    #[test]
    fn mlp_learns_xor() {
        // XOR is not linearly separable; a tiny hidden MLP can learn it.
        let examples: Vec<(Vec<f32>, usize)> = vec![
            (vec![0.0, 0.0], 0),
            (vec![0.0, 1.0], 1),
            (vec![1.0, 0.0], 1),
            (vec![1.0, 1.0], 0),
        ];

        let mut model = MlpClassifier::new(2, 8, 2, 7);
        let lr = 0.2;

        for _ in 0..2000 {
            for (x, y) in &examples {
                let grad = model.backward(x, *y);
                model.apply_sgd(&grad, lr);
            }
        }

        for (x, y) in &examples {
            assert_eq!(
                model.predict(x),
                *y,
                "failed XOR input {:?}, logits {:?}",
                x,
                model.forward(x).logits
            );
        }
    }

    #[test]
    fn flatten_and_load_round_trip() {
        let model = MlpClassifier::new(3, 4, 2, 13);
        let params = model.flatten_params();
        let mut restored = MlpClassifier::new(3, 4, 2, 99);
        restored.load_flat_params(&params);
        assert_eq!(model.w1, restored.w1);
        assert_eq!(model.b1, restored.b1);
        assert_eq!(model.w2, restored.w2);
        assert_eq!(model.b2, restored.b2);
    }

    #[test]
    fn backward_batch_matches_per_example_average() {
        let input_dim = 4;
        let hidden_dim = 5;
        let output_dim = 3;
        let model = MlpClassifier::new(input_dim, hidden_dim, output_dim, 21);
        let batch_size = 3;

        let inputs: Vec<f32> = (0..input_dim * batch_size)
            .map(|i| (i as f32) * 0.13 - 0.4)
            .collect();
        let targets = vec![0, 2, 1];

        let (grad_batch, _dx, loss_batch) = model.backward_batch(&inputs, &targets, batch_size);

        let mut dw1 = vec![0.0f32; grad_batch.dw1.len()];
        let mut db1 = vec![0.0f32; grad_batch.db1.len()];
        let mut dw2 = vec![0.0f32; grad_batch.dw2.len()];
        let mut db2 = vec![0.0f32; grad_batch.db2.len()];
        let mut loss_sum = 0.0f32;

        for b in 0..batch_size {
            let x = &inputs[b * input_dim..(b + 1) * input_dim];
            let g = model.backward(x, targets[b]);
            for (a, v) in dw1.iter_mut().zip(g.dw1.iter()) {
                *a += v;
            }
            for (a, v) in db1.iter_mut().zip(g.db1.iter()) {
                *a += v;
            }
            for (a, v) in dw2.iter_mut().zip(g.dw2.iter()) {
                *a += v;
            }
            for (a, v) in db2.iter_mut().zip(g.db2.iter()) {
                *a += v;
            }
            loss_sum += model.loss(x, targets[b]);
        }

        let inv = 1.0 / batch_size as f32;
        for v in dw1.iter_mut() {
            *v *= inv;
        }
        for v in db1.iter_mut() {
            *v *= inv;
        }
        for v in dw2.iter_mut() {
            *v *= inv;
        }
        for v in db2.iter_mut() {
            *v *= inv;
        }

        let tol = 1e-5;
        for (a, b) in grad_batch.dw1.iter().zip(dw1.iter()) {
            assert!((a - b).abs() < tol, "dw1 mismatch: {} vs {}", a, b);
        }
        for (a, b) in grad_batch.db1.iter().zip(db1.iter()) {
            assert!((a - b).abs() < tol, "db1 mismatch: {} vs {}", a, b);
        }
        for (a, b) in grad_batch.dw2.iter().zip(dw2.iter()) {
            assert!((a - b).abs() < tol, "dw2 mismatch: {} vs {}", a, b);
        }
        for (a, b) in grad_batch.db2.iter().zip(db2.iter()) {
            assert!((a - b).abs() < tol, "db2 mismatch: {} vs {}", a, b);
        }
        assert!((loss_batch - loss_sum * inv).abs() < tol);
    }

    #[test]
    fn embedding_mlp_learns_xor_with_continuous() {
        // Represent XOR inputs as two token IDs from a 2-token vocab, with no
        // continuous features.  The embedding layer must learn to distinguish
        // the two "bit" values.
        let examples: Vec<(u32, u32, usize)> = vec![(0, 0, 0), (0, 1, 1), (1, 0, 1), (1, 1, 0)];

        let mut model = EmbeddingMlpClassifier::new(2, 8, 2, 0, 16, 2, 7, None);
        let lr = 0.2;

        for _ in 0..8000 {
            for &(t1, t2, y) in &examples {
                let grad = model.backward_batch(&[t1, t2], None, &[y], 1);
                model.apply_sgd(&grad.0, lr);
            }
        }

        for &(t1, t2, y) in &examples {
            assert_eq!(
                model.predict(&[t1, t2], None),
                y,
                "failed XOR input ({}, {})",
                t1,
                t2
            );
        }
    }

    #[test]
    fn embedding_mlp_batch_matches_single_example() {
        let vocab_size = 5;
        let embed_dim = 3;
        let n_context = 2;
        let continuous_dim = 4;
        let model = EmbeddingMlpClassifier::new(
            vocab_size,
            embed_dim,
            n_context,
            continuous_dim,
            6,
            3,
            11,
            None,
        );

        let token_ids = vec![0u32, 2, 4, 1];
        let continuous: Vec<f32> = (0..(2 * continuous_dim) as i32)
            .map(|i| i as f32 * 0.1 - 0.5)
            .collect();
        let targets = vec![0, 2];

        let (batch_grad, batch_loss) =
            model.backward_batch(&token_ids, Some(&continuous), &targets, 2);

        // Compare against two single-example backward passes.
        let mut dembedding = vec![0.0f32; vocab_size * embed_dim];
        let mut dw1 = vec![0.0f32; model.mlp.w1.len()];
        let mut db1 = vec![0.0f32; model.mlp.b1.len()];
        let mut dw2 = vec![0.0f32; model.mlp.w2.len()];
        let mut db2 = vec![0.0f32; model.mlp.b2.len()];
        let mut loss_sum = 0.0f32;

        for b in 0..2 {
            let t = &token_ids[b * n_context..(b + 1) * n_context];
            let c = &continuous[b * continuous_dim..(b + 1) * continuous_dim];
            let (g, loss) = model.backward_batch(t, Some(c), &[targets[b]], 1);
            loss_sum += loss;

            for (a, v) in dembedding.iter_mut().zip(g.dembedding.iter()) {
                *a += v;
            }
            for (a, v) in dw1.iter_mut().zip(g.mlp.dw1.iter()) {
                *a += v;
            }
            for (a, v) in db1.iter_mut().zip(g.mlp.db1.iter()) {
                *a += v;
            }
            for (a, v) in dw2.iter_mut().zip(g.mlp.dw2.iter()) {
                *a += v;
            }
            for (a, v) in db2.iter_mut().zip(g.mlp.db2.iter()) {
                *a += v;
            }
        }

        let inv = 1.0 / 2.0;
        for v in dembedding.iter_mut() {
            *v *= inv;
        }
        for v in dw1.iter_mut() {
            *v *= inv;
        }
        for v in db1.iter_mut() {
            *v *= inv;
        }
        for v in dw2.iter_mut() {
            *v *= inv;
        }
        for v in db2.iter_mut() {
            *v *= inv;
        }

        let tol = 1e-5;
        for (a, b) in batch_grad.dembedding.iter().zip(dembedding.iter()) {
            assert!((a - b).abs() < tol, "dembedding mismatch: {} vs {}", a, b);
        }
        for (a, b) in batch_grad.mlp.dw1.iter().zip(dw1.iter()) {
            assert!((a - b).abs() < tol, "dw1 mismatch: {} vs {}", a, b);
        }
        for (a, b) in batch_grad.mlp.db1.iter().zip(db1.iter()) {
            assert!((a - b).abs() < tol, "db1 mismatch: {} vs {}", a, b);
        }
        for (a, b) in batch_grad.mlp.dw2.iter().zip(dw2.iter()) {
            assert!((a - b).abs() < tol, "dw2 mismatch: {} vs {}", a, b);
        }
        for (a, b) in batch_grad.mlp.db2.iter().zip(db2.iter()) {
            assert!((a - b).abs() < tol, "db2 mismatch: {} vs {}", a, b);
        }
        assert!((batch_loss - loss_sum * inv).abs() < tol);
    }

    #[test]
    fn rope_rotates_differently_by_position() {
        let mut v1 = vec![1.0f32, 0.0, 0.0, 1.0];
        let mut v2 = v1.clone();
        apply_rope_in_place(&mut v1, 4, 1, 10000.0);
        apply_rope_in_place(&mut v2, 4, 2, 10000.0);
        // Same vector rotated by different positions should differ.
        assert!(v1.iter().zip(v2.iter()).any(|(a, b)| (a - b).abs() > 1e-6));
    }

    #[test]
    fn rope_inverse_recovers_original() {
        let original = vec![0.3f32, -0.7, 1.2, 0.4];
        let mut rotated = original.clone();
        apply_rope_in_place(&mut rotated, 4, 3, 10000.0);
        let mut recovered = rotated.clone();
        apply_rope_inv_in_place(&mut recovered, 4, 3, 10000.0);

        let tol = 1e-5;
        for (a, b) in original.iter().zip(recovered.iter()) {
            assert!((a - b).abs() < tol, "RoPE inverse mismatch: {} vs {}", a, b);
        }
    }

    #[test]
    fn embedding_mlp_batch_with_rope_matches_per_example() {
        let vocab_size = 5;
        let embed_dim = 4;
        let n_context = 2;
        let continuous_dim = 0;
        let model = EmbeddingMlpClassifier::new(
            vocab_size,
            embed_dim,
            n_context,
            continuous_dim,
            6,
            3,
            11,
            Some(10000.0),
        );

        let token_ids = vec![0u32, 2, 4, 1];
        let targets = vec![0, 2];

        let (batch_grad, batch_loss) = model.backward_batch(&token_ids, None, &targets, 2);

        // Compare against two single-example backward passes.
        let mut dembedding = vec![0.0f32; vocab_size * embed_dim];
        let mut dw1 = vec![0.0f32; model.mlp.w1.len()];
        let mut db1 = vec![0.0f32; model.mlp.b1.len()];
        let mut dw2 = vec![0.0f32; model.mlp.w2.len()];
        let mut db2 = vec![0.0f32; model.mlp.b2.len()];
        let mut loss_sum = 0.0f32;

        for b in 0..2 {
            let t = &token_ids[b * n_context..(b + 1) * n_context];
            let (g, loss) = model.backward_batch(t, None, &[targets[b]], 1);
            loss_sum += loss;

            for (a, v) in dembedding.iter_mut().zip(g.dembedding.iter()) {
                *a += v;
            }
            for (a, v) in dw1.iter_mut().zip(g.mlp.dw1.iter()) {
                *a += v;
            }
            for (a, v) in db1.iter_mut().zip(g.mlp.db1.iter()) {
                *a += v;
            }
            for (a, v) in dw2.iter_mut().zip(g.mlp.dw2.iter()) {
                *a += v;
            }
            for (a, v) in db2.iter_mut().zip(g.mlp.db2.iter()) {
                *a += v;
            }
        }

        let inv = 1.0 / 2.0;
        for v in dembedding.iter_mut() {
            *v *= inv;
        }
        for v in dw1.iter_mut() {
            *v *= inv;
        }
        for v in db1.iter_mut() {
            *v *= inv;
        }
        for v in dw2.iter_mut() {
            *v *= inv;
        }
        for v in db2.iter_mut() {
            *v *= inv;
        }

        let tol = 1e-5;
        for (a, b) in batch_grad.dembedding.iter().zip(dembedding.iter()) {
            assert!((a - b).abs() < tol, "dembedding mismatch: {} vs {}", a, b);
        }
        for (a, b) in batch_grad.mlp.dw1.iter().zip(dw1.iter()) {
            assert!((a - b).abs() < tol, "dw1 mismatch: {} vs {}", a, b);
        }
        for (a, b) in batch_grad.mlp.db1.iter().zip(db1.iter()) {
            assert!((a - b).abs() < tol, "db1 mismatch: {} vs {}", a, b);
        }
        for (a, b) in batch_grad.mlp.dw2.iter().zip(dw2.iter()) {
            assert!((a - b).abs() < tol, "dw2 mismatch: {} vs {}", a, b);
        }
        for (a, b) in batch_grad.mlp.db2.iter().zip(db2.iter()) {
            assert!((a - b).abs() < tol, "db2 mismatch: {} vs {}", a, b);
        }
        assert!((batch_loss - loss_sum * inv).abs() < tol);
    }
}