ops: differentiable autograd nodes + per-op grad-check tests

ops.rs wraps each Tensor op as a Var node with its backward closure (forward
caches captured by move). swiglu = mul(silu(gate), up); attention is composed
(matmul+scale+softmax+matmul), no fused kernel. tests/autograd.rs grad-checks
every op via the L=sum(W∘out) template, plus a fan-out grad-accumulation test
(dL/dx=4x) and an end-to-end composed-attention grad-check (dQ/dK/dV). Adds
xtrain-cuda dev-dep for device selection in tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-autodiff/src/ops.rs

@@ -0,0 +1,172 @@
+//! 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);
+        }),
+    )
+}