xtrain/crates/xtrain-autodiff/src/ops.rs

//! Differentiable ops as autograd nodes (Phase T4).
//!
//! Each function runs the forward [`Tensor`] kernel, then builds a [`Var`] whose
//! backward closure computes the analytic gradient (see
//! `docs/03-autograd-engine.md` for the math) and pushes it to each parent via
//! [`Var::push_grad`] (which SUMs — correct under fan-out). Forward outputs that
//! the backward needs (softmax `y`, rms `inv_rms`, cross-entropy `probs`) are
//! cached by moving them into the closure.
//!
//! Attention is NOT a node here: it is composed from `matmul` + `scale` +
//! `softmax` in user code, and its backward falls out of theirs.

#![cfg(not(no_cuda))]

use crate::tape::Var;
use xtrain_tensor::Tensor;

/// `C = A @ B` (2D). Backward: `dA = dC @ Bᵀ`, `dB = Aᵀ @ dC`.
pub fn matmul(a: &Var, b: &Var) -> Var {
    let out = a.value().matmul(&b.value());
    Var::from_op(
        out,
        vec![a.clone(), b.clone()],
        Box::new(|dc, parents| {
            let a = parents[0].value();
            let b = parents[1].value();
            let (da, db) = Tensor::matmul_backward(&a, &b, dc);
            Var::push_grad(&parents[0], da);
            Var::push_grad(&parents[1], db);
        }),
    )
}

/// Elementwise `out = a + b` (same shape). Backward: grad flows unchanged to both.
pub fn add(a: &Var, b: &Var) -> Var {
    let out = a.value().add(&b.value());
    Var::from_op(
        out,
        vec![a.clone(), b.clone()],
        Box::new(|d, parents| {
            Var::push_grad(&parents[0], d.clone());
            Var::push_grad(&parents[1], d.clone());
        }),
    )
}

/// Elementwise `out = a * b` (Hadamard). Backward: `da = d∘b`, `db = d∘a`.
pub fn mul(a: &Var, b: &Var) -> Var {
    let out = a.value().mul(&b.value());
    Var::from_op(
        out,
        vec![a.clone(), b.clone()],
        Box::new(|d, parents| {
            let a = parents[0].value();
            let b = parents[1].value();
            Var::push_grad(&parents[0], d.mul(&b));
            Var::push_grad(&parents[1], d.mul(&a));
        }),
    )
}

/// Broadcast bias add: `out[r,c] = x[r,c] + bias[c]`. Backward: `dx = d`,
/// `dbias[c] = sum_r d[r,c]` (sum over the broadcast dim).
pub fn add_bias(x: &Var, bias: &Var) -> Var {
    let out = x.value().add_bias(&bias.value());
    Var::from_op(
        out,
        vec![x.clone(), bias.clone()],
        Box::new(|d, parents| {
            Var::push_grad(&parents[0], d.clone());
            Var::push_grad(&parents[1], d.sum_rows());
        }),
    )
}

/// Scale by a constant: `out = x * alpha`. Backward: `dx = d * alpha`.
pub fn scale(x: &Var, alpha: f32) -> Var {
    let out = x.value().scale(alpha);
    Var::from_op(
        out,
        vec![x.clone()],
        Box::new(move |d, parents| {
            Var::push_grad(&parents[0], d.scale(alpha));
        }),
    )
}

/// RMSNorm: `y = x * rsqrt(mean(x²)+eps) * gamma`. Caches `inv_rms` for backward.
pub fn rms_norm(x: &Var, gamma: &Var, eps: f32) -> Var {
    let (y, inv_rms) = x.value().rms_norm(&gamma.value(), eps);
    Var::from_op(
        y,
        vec![x.clone(), gamma.clone()],
        Box::new(move |dy, parents| {
            let x = parents[0].value();
            let gamma = parents[1].value();
            let (dx, dgamma) = Tensor::rms_norm_backward(&x, &gamma, dy, &inv_rms);
            Var::push_grad(&parents[0], dx);
            Var::push_grad(&parents[1], dgamma);
        }),
    )
}

/// SiLU: `y = x * sigmoid(x)`. Backward uses the forward `x`.
pub fn silu(x: &Var) -> Var {
    let out = x.value().silu();
    Var::from_op(
        out,
        vec![x.clone()],
        Box::new(|dy, parents| {
            let x = parents[0].value();
            Var::push_grad(&parents[0], Tensor::silu_backward(&x, dy));
        }),
    )
}

/// SwiGLU (SiLU-gated GLU): `out = silu(gate) ∘ up`. Composed from `silu` + `mul`
/// so its backward comes from theirs — no dedicated kernel needed.
pub fn swiglu(gate: &Var, up: &Var) -> Var {
    mul(&silu(gate), up)
}

/// RoPE (rotate_half) over `x:[tokens,heads,head_dim]`. Orthogonal map, so the
/// backward is the inverse rotation of `dy` — no cached forward values needed.
pub fn rope(x: &Var, theta: f32) -> Var {
    let out = x.value().rope(theta);
    Var::from_op(
        out,
        vec![x.clone()],
        Box::new(move |dy, parents| {
            Var::push_grad(&parents[0], Tensor::rope_backward(dy, theta));
        }),
    )
}

/// Row-wise softmax. Caches the output `y` for the Jacobian backward.
pub fn softmax(x: &Var) -> Var {
    let y = x.value().softmax();
    let y_cache = y.clone();
    Var::from_op(
        y,
        vec![x.clone()],
        Box::new(move |dy, parents| {
            Var::push_grad(&parents[0], Tensor::softmax_backward(&y_cache, dy));
        }),
    )
}

/// Cross-entropy mean loss over logits `x:[rows,cols]` with one I32 target per
/// row. Returns a scalar [`Var`]. Backward: `dx = (probs - onehot)/rows`,
/// scaled by the upstream scalar grad.
pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
    let (probs, per_row) = x.value().cross_entropy(target);
    let rows = x.value().shape()[0];
    // Mean loss as a host scalar wrapped back into a [1] tensor.
    let mean = per_row.to_device(xtrain_tensor::Device::Cpu);
    let mean_val: f32 = mean.as_slice::<f32>().iter().sum::<f32>() / rows as f32;
    let loss = Tensor::from_slice(&[mean_val], &[1]).to_device(x.value().device());

    let target = target.clone();
    Var::from_op(
        loss,
        vec![x.clone()],
        Box::new(move |d, parents| {
            // `d` is the scalar upstream grad (1.0 when this is the loss root).
            let upstream = d.to_device(xtrain_tensor::Device::Cpu).as_slice::<f32>()[0];
            let scale = upstream / rows as f32;
            let dx = Tensor::cross_entropy_backward(&probs, &target, scale);
            Var::push_grad(&parents[0], dx);
        }),
    )
}