Merge pull request 'rtx-autograd: make the tape correct + sound on real backends' (#3) from autograd-trainable-real-backend into main

crates/core/rtx-autograd/Cargo.toml

@@ -36,6 +36,10 @@ criterion.workspace = true
 approx = "0.5"
 rand = "0.8"
 
+# Real CPU backend for numerical gradient-checking the tape (no cycle:
+# rtx-backend-cpu depends only on rtx-backend / rtx-tensor, not rtx-autograd).
+rtx-backend-cpu = { path = "../rtx-backend-cpu" }
+
 [features]
 default = []
 disabled_tests = []

crates/core/rtx-autograd/src/autodiff/backend.rs

@@ -20,6 +20,7 @@ use rtx_backend::{AutodiffBackend, Backend, DeviceId, DeviceOps, GradientMap};
 use super::context::is_grad_enabled;
 use super::graph::{GradientStorage, backward_impl};
 use super::node::{AutodiffNode, GradTensor, ParentRef};
+use super::ops::into_dim;
 use super::tensor::AutodiffTensor;
 
 // ============================================================================
@@ -1303,12 +1304,12 @@ where
         // 2. The match arm ensures D matches the target dimension (e.g., D=1 matches D1)
         // 3. `ones` is a valid B::TensorPrimitive<D> and we copy it to the same type with explicit D
         let initial_grad = match D {
-            1 => GradTensor::D1(unsafe { std::mem::transmute_copy(&ones) }),
+            1 => GradTensor::D1(into_dim::<B, D, 1>(ones)),
-            2 => GradTensor::D2(unsafe { std::mem::transmute_copy(&ones) }),
+            2 => GradTensor::D2(into_dim::<B, D, 2>(ones)),
-            3 => GradTensor::D3(unsafe { std::mem::transmute_copy(&ones) }),
+            3 => GradTensor::D3(into_dim::<B, D, 3>(ones)),
-            4 => GradTensor::D4(unsafe { std::mem::transmute_copy(&ones) }),
+            4 => GradTensor::D4(into_dim::<B, D, 4>(ones)),
-            5 => GradTensor::D5(unsafe { std::mem::transmute_copy(&ones) }),
+            5 => GradTensor::D5(into_dim::<B, D, 5>(ones)),
-            6 => GradTensor::D6(unsafe { std::mem::transmute_copy(&ones) }),
+            6 => GradTensor::D6(into_dim::<B, D, 6>(ones)),
             _ => {
                 return Err(AutogradError::UnsupportedDimension(D).into());
             }

crates/core/rtx-autograd/src/autodiff/graph.rs

@@ -212,19 +212,24 @@ where
         b.ndim()
     );
 
-    // For now, just return b (proper accumulation requires Backend::add)
+    // When a tensor fans out (is used by more than one downstream op, or
-    // This will be implemented properly when we have access to Backend operations
+    // appears more than once as a parent), each path contributes a gradient
-    match (&a, &b) {
+    // and the contributions must be *summed* at the shared sink. The previous
-        (GradTensor::D1(_), GradTensor::D1(_)) => Ok(b),
+    // body returned `b` and discarded `a`, so any reuse (residuals, `mul(s, s)`,
-        (GradTensor::D2(_), GradTensor::D2(_)) => Ok(b),
+    // shared Q/K/V inputs — universal in transformers) silently dropped a path
-        (GradTensor::D3(_), GradTensor::D3(_)) => Ok(b),
+    // and produced gradients that were a fraction of the true value. Sum them
-        (GradTensor::D4(_), GradTensor::D4(_)) => Ok(b),
+    // with the backend's elementwise add.
-        (GradTensor::D5(_), GradTensor::D5(_)) => Ok(b),
+    let ndim_a = a.ndim();
-        (GradTensor::D6(_), GradTensor::D6(_)) => Ok(b),
+    let ndim_b = b.ndim();
+    match (a, b) {
+        (GradTensor::D1(x), GradTensor::D1(y)) => Ok(GradTensor::D1(B::add(x, y))),
+        (GradTensor::D2(x), GradTensor::D2(y)) => Ok(GradTensor::D2(B::add(x, y))),
+        (GradTensor::D3(x), GradTensor::D3(y)) => Ok(GradTensor::D3(B::add(x, y))),
+        (GradTensor::D4(x), GradTensor::D4(y)) => Ok(GradTensor::D4(B::add(x, y))),
+        (GradTensor::D5(x), GradTensor::D5(y)) => Ok(GradTensor::D5(B::add(x, y))),
+        (GradTensor::D6(x), GradTensor::D6(y)) => Ok(GradTensor::D6(B::add(x, y))),
         _ => Err(AutogradError::DimensionMismatch(format!(
-            "Cannot accumulate gradients of different dimensions: {} and {}",
+            "Cannot accumulate gradients of different dimensions: {ndim_a} and {ndim_b}"
-            a.ndim(),
-            b.ndim()
         ))),
     }
 }

crates/core/rtx-autograd/src/autodiff/ops/activations.rs

@@ -9,6 +9,7 @@
 //!                = `grad_output * sigmoid(x) * (1 + x * (1 - sigmoid(x)))`
 
 use crate::autodiff::node::{AutodiffBackwardFn, GradTensor, SavedTensor};
+use crate::autodiff::ops::into_dim;
 use crate::error::{AutogradError, Result};
 use rtx_backend::Backend;
 use std::marker::PhantomData;
@@ -597,32 +598,32 @@ where
     match (grad, D) {
         (GradTensor::D1(g), 1) => {
             // SAFETY: D=1 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<1>>(derivative) };
+            let t = into_dim::<B, D, 1>(derivative.clone());
             Ok(GradTensor::D1(B::mul(g.clone(), t)))
         }
         (GradTensor::D2(g), 2) => {
             // SAFETY: D=2 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<2>>(derivative) };
+            let t = into_dim::<B, D, 2>(derivative.clone());
             Ok(GradTensor::D2(B::mul(g.clone(), t)))
         }
         (GradTensor::D3(g), 3) => {
             // SAFETY: D=3 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<3>>(derivative) };
+            let t = into_dim::<B, D, 3>(derivative.clone());
             Ok(GradTensor::D3(B::mul(g.clone(), t)))
         }
         (GradTensor::D4(g), 4) => {
             // SAFETY: D=4 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<4>>(derivative) };
+            let t = into_dim::<B, D, 4>(derivative.clone());
             Ok(GradTensor::D4(B::mul(g.clone(), t)))
         }
         (GradTensor::D5(g), 5) => {
             // SAFETY: D=5 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<5>>(derivative) };
+            let t = into_dim::<B, D, 5>(derivative.clone());
             Ok(GradTensor::D5(B::mul(g.clone(), t)))
         }
         (GradTensor::D6(g), 6) => {
             // SAFETY: D=6 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<6>>(derivative) };
+            let t = into_dim::<B, D, 6>(derivative.clone());
             Ok(GradTensor::D6(B::mul(g.clone(), t)))
         }
         _ => Err(AutogradError::DimensionMismatch(

crates/core/rtx-autograd/src/autodiff/ops/basic.rs

@@ -9,6 +9,7 @@
 //! - Neg: `grad_input = -grad_out`
 
 use crate::autodiff::node::{AutodiffBackwardFn, GradTensor, SavedTensor};
+use crate::autodiff::ops::into_dim;
 use crate::error::{AutogradError, Result};
 use rtx_backend::Backend;
 use std::marker::PhantomData;
@@ -294,32 +295,32 @@ where
     match (grad, D) {
         (GradTensor::D1(g), 1) => {
             // SAFETY: D=1 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<1>>(tensor) };
+            let t = into_dim::<B, D, 1>(tensor.clone());
             Ok(GradTensor::D1(B::mul(g, t)))
         }
         (GradTensor::D2(g), 2) => {
             // SAFETY: D=2 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<2>>(tensor) };
+            let t = into_dim::<B, D, 2>(tensor.clone());
             Ok(GradTensor::D2(B::mul(g, t)))
         }
         (GradTensor::D3(g), 3) => {
             // SAFETY: D=3 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<3>>(tensor) };
+            let t = into_dim::<B, D, 3>(tensor.clone());
             Ok(GradTensor::D3(B::mul(g, t)))
         }
         (GradTensor::D4(g), 4) => {
             // SAFETY: D=4 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<4>>(tensor) };
+            let t = into_dim::<B, D, 4>(tensor.clone());
             Ok(GradTensor::D4(B::mul(g, t)))
         }
         (GradTensor::D5(g), 5) => {
             // SAFETY: D=5 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<5>>(tensor) };
+            let t = into_dim::<B, D, 5>(tensor.clone());
             Ok(GradTensor::D5(B::mul(g, t)))
         }
         (GradTensor::D6(g), 6) => {
             // SAFETY: D=6 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<6>>(tensor) };
+            let t = into_dim::<B, D, 6>(tensor.clone());
             Ok(GradTensor::D6(B::mul(g, t)))
         }
         _ => Err(AutogradError::DimensionMismatch(
@@ -345,32 +346,32 @@ where
     match (grad, D) {
         (GradTensor::D1(g), 1) => {
             // SAFETY: D=1 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<1>>(tensor) };
+            let t = into_dim::<B, D, 1>(tensor.clone());
             Ok(GradTensor::D1(B::div(g, t)))
         }
         (GradTensor::D2(g), 2) => {
             // SAFETY: D=2 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<2>>(tensor) };
+            let t = into_dim::<B, D, 2>(tensor.clone());
             Ok(GradTensor::D2(B::div(g, t)))
         }
         (GradTensor::D3(g), 3) => {
             // SAFETY: D=3 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<3>>(tensor) };
+            let t = into_dim::<B, D, 3>(tensor.clone());
             Ok(GradTensor::D3(B::div(g, t)))
         }
         (GradTensor::D4(g), 4) => {
             // SAFETY: D=4 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<4>>(tensor) };
+            let t = into_dim::<B, D, 4>(tensor.clone());
             Ok(GradTensor::D4(B::div(g, t)))
         }
         (GradTensor::D5(g), 5) => {
             // SAFETY: D=5 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<5>>(tensor) };
+            let t = into_dim::<B, D, 5>(tensor.clone());
             Ok(GradTensor::D5(B::div(g, t)))
         }
         (GradTensor::D6(g), 6) => {
             // SAFETY: D=6 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<6>>(tensor) };
+            let t = into_dim::<B, D, 6>(tensor.clone());
             Ok(GradTensor::D6(B::div(g, t)))
         }
         _ => Err(AutogradError::DimensionMismatch(

crates/core/rtx-autograd/src/autodiff/ops/llm.rs

@@ -8,6 +8,7 @@
 //! - FlashAttention: Fused attention mechanism
 
 use crate::autodiff::node::{AutodiffBackwardFn, GradTensor, SavedTensor};
+use crate::autodiff::ops::into_dim;
 use crate::error::{AutogradError, Result};
 use rtx_backend::Backend;
 use std::marker::PhantomData;
@@ -66,11 +67,16 @@ where
         let grad_out_tensor = extract_tensor::<B, D>(&grad_output)?;
         let grad_times_soft = B::mul(grad_out_tensor.clone(), softmax_out.clone());
 
-        // Step 2: sum along dim (keeping dims for broadcast)
+        // Step 2: sum along dim. `sum_dim` is keep-dim, so this is the full
+        // shape with `dim` collapsed to 1.
         let sum_grad_soft = B::sum_dim(grad_times_soft, *dim);
 
-        // Step 3: grad_y - sum
+        // Step 3: grad_y - sum. The elementwise backends require identical
-        let grad_minus_sum = B::sub(grad_out_tensor, sum_grad_soft);
+        // shapes (no implicit broadcasting), so tile the per-row sum back to
+        // the full width along `dim` before subtracting.
+        let full = B::shape(softmax_out)[*dim];
+        let sum_broadcast = broadcast_along_dim::<B, D>(sum_grad_soft, *dim, full);
+        let grad_minus_sum = B::sub(grad_out_tensor, sum_broadcast);
 
         // Step 4: softmax * (grad_y - sum)
         let grad_input = B::mul(softmax_out.clone(), grad_minus_sum);
@@ -595,6 +601,37 @@ where
 ///
 /// The Backend trait's softmax already implements this internally.
 
+/// Tile a keep-dim–reduced tensor (size 1 along `dim`) back to `full`
+/// elements along `dim`.
+///
+/// The elementwise backends (`add`/`sub`/`mul`) assert identical operand
+/// shapes — there is no implicit broadcasting — so a per-row reduction must
+/// be materialised to the full shape before it can be combined with the
+/// unreduced tensor. Done generically via host data so it works for any
+/// backend and dimension.
+fn broadcast_along_dim<B: Backend, const D: usize>(
+    reduced: B::TensorPrimitive<D>,
+    dim: usize,
+    full: usize,
+) -> B::TensorPrimitive<D> {
+    let shape = B::shape(&reduced);
+    let device = B::device(&reduced);
+    let data = B::to_data(&reduced);
+    let outer: usize = shape[..dim].iter().product();
+    let inner: usize = shape[dim + 1..].iter().product();
+    let mut out_shape = shape;
+    out_shape[dim] = full;
+    let mut out = Vec::with_capacity(outer * full * inner);
+    for o in 0..outer {
+        for _f in 0..full {
+            for i in 0..inner {
+                out.push(data[o * inner + i]);
+            }
+        }
+    }
+    B::from_data(&out, out_shape, &device)
+}
+
 /// Numerically stable softmax backward using log-sum-exp values.
 ///
 /// Given:
@@ -616,9 +653,13 @@ fn stable_softmax_backward<B: Backend>(
 where
     B::TensorPrimitive<3>: Clone,
 {
-    // Standard softmax backward: grad_x = y * (grad_y - sum(grad_y * y, dim))
+    // Standard softmax backward: grad_x = y * (grad_y - sum(grad_y * y, dim)).
+    // `sum_dim` is keep-dim; tile the per-row sum back to full width along
+    // `dim` before subtracting, since the elementwise backends don't broadcast.
     let grad_times_softmax = B::mul(grad_output.clone(), softmax_output.clone());
     let rowsum = B::sum_dim(grad_times_softmax, dim);
+    let full = B::shape(&softmax_output)[dim];
+    let rowsum = broadcast_along_dim::<B, 3>(rowsum, dim, full);
     let grad_minus_rowsum = B::sub(grad_output, rowsum);
     let grad_input = B::mul(softmax_output, grad_minus_rowsum);
 
@@ -634,16 +675,22 @@ where
 /// 3. t is a valid reference and we create a copy to the correctly-typed primitive
 fn extract_tensor<B: Backend, const D: usize>(grad: &GradTensor<B>) -> Result<B::TensorPrimitive<D>>
 where
-    B::TensorPrimitive<D>: Clone,
+    B::TensorPrimitive<1>: Clone,
+    B::TensorPrimitive<2>: Clone,
+    B::TensorPrimitive<3>: Clone,
+    B::TensorPrimitive<4>: Clone,
+    B::TensorPrimitive<5>: Clone,
+    B::TensorPrimitive<6>: Clone,
 {
-    // SAFETY: D is verified by each match arm; layout identical across D values
+    // `into_dim` consumes the cloned variant tensor (transmute_copy +
+    // mem::forget) and reinterprets it at the runtime dimension `D`.
     match (grad, D) {
-        (GradTensor::D1(t), 1) => Ok(unsafe { std::mem::transmute_copy(t) }),
+        (GradTensor::D1(t), 1) => Ok(into_dim::<B, 1, D>(t.clone())),
-        (GradTensor::D2(t), 2) => Ok(unsafe { std::mem::transmute_copy(t) }),
+        (GradTensor::D2(t), 2) => Ok(into_dim::<B, 2, D>(t.clone())),
-        (GradTensor::D3(t), 3) => Ok(unsafe { std::mem::transmute_copy(t) }),
+        (GradTensor::D3(t), 3) => Ok(into_dim::<B, 3, D>(t.clone())),
-        (GradTensor::D4(t), 4) => Ok(unsafe { std::mem::transmute_copy(t) }),
+        (GradTensor::D4(t), 4) => Ok(into_dim::<B, 4, D>(t.clone())),
-        (GradTensor::D5(t), 5) => Ok(unsafe { std::mem::transmute_copy(t) }),
+        (GradTensor::D5(t), 5) => Ok(into_dim::<B, 5, D>(t.clone())),
-        (GradTensor::D6(t), 6) => Ok(unsafe { std::mem::transmute_copy(t) }),
+        (GradTensor::D6(t), 6) => Ok(into_dim::<B, 6, D>(t.clone())),
         _ => Err(AutogradError::DimensionMismatch(
             "llm operation".to_string(),
         )),
@@ -660,12 +707,12 @@ where
 fn wrap_tensor<B: Backend, const D: usize>(tensor: B::TensorPrimitive<D>) -> Result<GradTensor<B>> {
     // SAFETY: D is verified by each match arm; layout identical across D values
     match D {
-        1 => Ok(GradTensor::D1(unsafe { std::mem::transmute_copy(&tensor) })),
+        1 => Ok(GradTensor::D1(into_dim::<B, D, 1>(tensor))),
-        2 => Ok(GradTensor::D2(unsafe { std::mem::transmute_copy(&tensor) })),
+        2 => Ok(GradTensor::D2(into_dim::<B, D, 2>(tensor))),
-        3 => Ok(GradTensor::D3(unsafe { std::mem::transmute_copy(&tensor) })),
+        3 => Ok(GradTensor::D3(into_dim::<B, D, 3>(tensor))),
-        4 => Ok(GradTensor::D4(unsafe { std::mem::transmute_copy(&tensor) })),
+        4 => Ok(GradTensor::D4(into_dim::<B, D, 4>(tensor))),
-        5 => Ok(GradTensor::D5(unsafe { std::mem::transmute_copy(&tensor) })),
+        5 => Ok(GradTensor::D5(into_dim::<B, D, 5>(tensor))),
-        6 => Ok(GradTensor::D6(unsafe { std::mem::transmute_copy(&tensor) })),
+        6 => Ok(GradTensor::D6(into_dim::<B, D, 6>(tensor))),
         _ => Err(AutogradError::UnsupportedDimension(D)),
     }
 }

crates/core/rtx-autograd/src/autodiff/ops/mod.rs

@@ -23,6 +23,38 @@ mod reductions;
 mod shapes;
 mod unary;
 
+use rtx_backend::Backend;
+
+/// Reinterpret an owned `TensorPrimitive<FROM>` as a `TensorPrimitive<TO>`.
+///
+/// The backward functions are generic over a const dimension `D` but must
+/// hand back a `GradTensor` whose variant matches `D` at runtime. Every call
+/// site selects, via a `match` on the runtime dimension, the arm whose `TO`
+/// equals the actual `D` — so in practice `FROM == TO` whenever this runs.
+///
+/// Implemented with **no `unsafe`**: it round-trips through the backend's safe
+/// data API (`to_data` → `from_data`) and rebuilds the `[usize; TO]` shape
+/// element-by-element. The previous implementation cast across the const
+/// generic with `transmute_copy`, which (a) double-freed the data buffer on
+/// any heap-backed backend — the source was left to drop alongside the
+/// bit-copy — and (b) was Undefined Behaviour under Stacked Borrows (a typed
+/// read punned through a reference of a different type). Both are gone here;
+/// the cost is one data copy per cast, which is negligible relative to the op
+/// itself and only happens on the gradient path.
+///
+/// For the never-executed arms where `FROM != TO`, the shape is padded with
+/// zeros / truncated; the value is discarded before use, so this is harmless.
+#[inline]
+pub(crate) fn into_dim<B: Backend, const FROM: usize, const TO: usize>(
+    src: B::TensorPrimitive<FROM>,
+) -> B::TensorPrimitive<TO> {
+    let from_shape = B::shape(&src);
+    let to_shape: [usize; TO] = std::array::from_fn(|i| if i < FROM { from_shape[i] } else { 0 });
+    let device = B::device(&src);
+    let data = B::to_data(&src);
+    B::from_data(&data, to_shape, &device)
+}
+
 // Re-export all backward functions
 pub use activations::{
     EluBackward, GeluBackward, LeakyReluBackward, ReluBackward, SigmoidBackward, SiluBackward,

crates/core/rtx-autograd/src/autodiff/ops/reductions.rs

@@ -15,6 +15,7 @@
 //! 3. All transmute_copy operations create owned copies from valid references/values
 
 use crate::autodiff::node::{AutodiffBackwardFn, GradTensor, SavedTensor};
+use crate::autodiff::ops::into_dim;
 use crate::error::{AutogradError, Result};
 use rtx_backend::Backend;
 use std::marker::PhantomData;
@@ -92,16 +93,20 @@ where
             _ => unreachable!(),
         };
 
-        // Create full tensor with the scalar value
+        // Create full tensor with the scalar value, then reinterpret its
+        // owned storage as the runtime dimension `D`. `into_dim` CONSUMES
+        // `grad_input` (transmute_copy + mem::forget), so the underlying
+        // `Vec` is freed exactly once — the previous `transmute_copy(&owned)`
+        // left both `grad_input` and the copy owning the same buffer, which
+        // double-freed under any real (heap-backed) backend.
         let grad_input = B::full(*shape, grad_scalar[0], &device);
-        // SAFETY: See module-level documentation. D is verified by each match arm.
         let grad_tensor = match D {
-            1 => GradTensor::D1(unsafe { std::mem::transmute_copy(&grad_input) }),
+            1 => GradTensor::D1(into_dim::<B, D, 1>(grad_input)),
-            2 => GradTensor::D2(unsafe { std::mem::transmute_copy(&grad_input) }),
+            2 => GradTensor::D2(into_dim::<B, D, 2>(grad_input)),
-            3 => GradTensor::D3(unsafe { std::mem::transmute_copy(&grad_input) }),
+            3 => GradTensor::D3(into_dim::<B, D, 3>(grad_input)),
-            4 => GradTensor::D4(unsafe { std::mem::transmute_copy(&grad_input) }),
+            4 => GradTensor::D4(into_dim::<B, D, 4>(grad_input)),
-            5 => GradTensor::D5(unsafe { std::mem::transmute_copy(&grad_input) }),
+            5 => GradTensor::D5(into_dim::<B, D, 5>(grad_input)),
-            6 => GradTensor::D6(unsafe { std::mem::transmute_copy(&grad_input) }),
+            6 => GradTensor::D6(into_dim::<B, D, 6>(grad_input)),
             _ => return Err(AutogradError::UnsupportedDimension(0)),
         };
 
@@ -163,12 +168,12 @@ where
 
         // SAFETY: See module-level documentation. D is verified by each match arm.
         let grad_input = match D {
-            1 => GradTensor::D1(unsafe { std::mem::transmute_copy(&ones) }),
+            1 => GradTensor::D1(into_dim::<B, D, 1>(ones)),
-            2 => GradTensor::D2(unsafe { std::mem::transmute_copy(&ones) }),
+            2 => GradTensor::D2(into_dim::<B, D, 2>(ones)),
-            3 => GradTensor::D3(unsafe { std::mem::transmute_copy(&ones) }),
+            3 => GradTensor::D3(into_dim::<B, D, 3>(ones)),
-            4 => GradTensor::D4(unsafe { std::mem::transmute_copy(&ones) }),
+            4 => GradTensor::D4(into_dim::<B, D, 4>(ones)),
-            5 => GradTensor::D5(unsafe { std::mem::transmute_copy(&ones) }),
+            5 => GradTensor::D5(into_dim::<B, D, 5>(ones)),
-            6 => GradTensor::D6(unsafe { std::mem::transmute_copy(&ones) }),
+            6 => GradTensor::D6(into_dim::<B, D, 6>(ones)),
             _ => return Err(AutogradError::UnsupportedDimension(0)),
         };
 
@@ -256,12 +261,12 @@ where
 
         // SAFETY: See module-level documentation. D is verified by each match arm.
         let grad_tensor = match D {
-            1 => GradTensor::D1(unsafe { std::mem::transmute_copy(&grad_input) }),
+            1 => GradTensor::D1(into_dim::<B, D, 1>(grad_input)),
-            2 => GradTensor::D2(unsafe { std::mem::transmute_copy(&grad_input) }),
+            2 => GradTensor::D2(into_dim::<B, D, 2>(grad_input)),
-            3 => GradTensor::D3(unsafe { std::mem::transmute_copy(&grad_input) }),
+            3 => GradTensor::D3(into_dim::<B, D, 3>(grad_input)),
-            4 => GradTensor::D4(unsafe { std::mem::transmute_copy(&grad_input) }),
+            4 => GradTensor::D4(into_dim::<B, D, 4>(grad_input)),
-            5 => GradTensor::D5(unsafe { std::mem::transmute_copy(&grad_input) }),
+            5 => GradTensor::D5(into_dim::<B, D, 5>(grad_input)),
-            6 => GradTensor::D6(unsafe { std::mem::transmute_copy(&grad_input) }),
+            6 => GradTensor::D6(into_dim::<B, D, 6>(grad_input)),
             _ => return Err(AutogradError::UnsupportedDimension(0)),
         };
 
@@ -412,12 +417,12 @@ where
 
         // SAFETY: See module-level documentation. D is verified by each match arm.
         let result = match D {
-            1 => GradTensor::D1(unsafe { std::mem::transmute_copy(&grad_input) }),
+            1 => GradTensor::D1(into_dim::<B, D, 1>(grad_input)),
-            2 => GradTensor::D2(unsafe { std::mem::transmute_copy(&grad_input) }),
+            2 => GradTensor::D2(into_dim::<B, D, 2>(grad_input)),
-            3 => GradTensor::D3(unsafe { std::mem::transmute_copy(&grad_input) }),
+            3 => GradTensor::D3(into_dim::<B, D, 3>(grad_input)),
-            4 => GradTensor::D4(unsafe { std::mem::transmute_copy(&grad_input) }),
+            4 => GradTensor::D4(into_dim::<B, D, 4>(grad_input)),
-            5 => GradTensor::D5(unsafe { std::mem::transmute_copy(&grad_input) }),
+            5 => GradTensor::D5(into_dim::<B, D, 5>(grad_input)),
-            6 => GradTensor::D6(unsafe { std::mem::transmute_copy(&grad_input) }),
+            6 => GradTensor::D6(into_dim::<B, D, 6>(grad_input)),
             _ => return Err(AutogradError::UnsupportedDimension(0)),
         };
 
@@ -505,12 +510,12 @@ where
 
         // SAFETY: See module-level documentation. D is verified by each match arm.
         let result = match D {
-            1 => GradTensor::D1(unsafe { std::mem::transmute_copy(&grad_input) }),
+            1 => GradTensor::D1(into_dim::<B, D, 1>(grad_input)),
-            2 => GradTensor::D2(unsafe { std::mem::transmute_copy(&grad_input) }),
+            2 => GradTensor::D2(into_dim::<B, D, 2>(grad_input)),
-            3 => GradTensor::D3(unsafe { std::mem::transmute_copy(&grad_input) }),
+            3 => GradTensor::D3(into_dim::<B, D, 3>(grad_input)),
-            4 => GradTensor::D4(unsafe { std::mem::transmute_copy(&grad_input) }),
+            4 => GradTensor::D4(into_dim::<B, D, 4>(grad_input)),
-            5 => GradTensor::D5(unsafe { std::mem::transmute_copy(&grad_input) }),
+            5 => GradTensor::D5(into_dim::<B, D, 5>(grad_input)),
-            6 => GradTensor::D6(unsafe { std::mem::transmute_copy(&grad_input) }),
+            6 => GradTensor::D6(into_dim::<B, D, 6>(grad_input)),
             _ => return Err(AutogradError::UnsupportedDimension(0)),
         };
 
@@ -555,27 +560,27 @@ where
     // SAFETY: See module-level documentation. D is verified by each match arm.
     match (grad, D) {
         (GradTensor::D1(g), 1) => {
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<1>>(tensor) };
+            let t = into_dim::<B, D, 1>(tensor.clone());
             Ok(GradTensor::D1(B::mul(g.clone(), t)))
         }
         (GradTensor::D2(g), 2) => {
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<2>>(tensor) };
+            let t = into_dim::<B, D, 2>(tensor.clone());
             Ok(GradTensor::D2(B::mul(g.clone(), t)))
         }
         (GradTensor::D3(g), 3) => {
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<3>>(tensor) };
+            let t = into_dim::<B, D, 3>(tensor.clone());
             Ok(GradTensor::D3(B::mul(g.clone(), t)))
         }
         (GradTensor::D4(g), 4) => {
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<4>>(tensor) };
+            let t = into_dim::<B, D, 4>(tensor.clone());
             Ok(GradTensor::D4(B::mul(g.clone(), t)))
         }
         (GradTensor::D5(g), 5) => {
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<5>>(tensor) };
+            let t = into_dim::<B, D, 5>(tensor.clone());
             Ok(GradTensor::D5(B::mul(g.clone(), t)))
         }
         (GradTensor::D6(g), 6) => {
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<6>>(tensor) };
+            let t = into_dim::<B, D, 6>(tensor.clone());
             Ok(GradTensor::D6(B::mul(g.clone(), t)))
         }
         _ => Err(AutogradError::DimensionMismatch(
@@ -723,12 +728,12 @@ where
 
         // SAFETY: See module-level documentation. D is verified by each match arm.
         let result = match D {
-            1 => GradTensor::D1(unsafe { std::mem::transmute_copy(&grad_input) }),
+            1 => GradTensor::D1(into_dim::<B, D, 1>(grad_input)),
-            2 => GradTensor::D2(unsafe { std::mem::transmute_copy(&grad_input) }),
+            2 => GradTensor::D2(into_dim::<B, D, 2>(grad_input)),
-            3 => GradTensor::D3(unsafe { std::mem::transmute_copy(&grad_input) }),
+            3 => GradTensor::D3(into_dim::<B, D, 3>(grad_input)),
-            4 => GradTensor::D4(unsafe { std::mem::transmute_copy(&grad_input) }),
+            4 => GradTensor::D4(into_dim::<B, D, 4>(grad_input)),
-            5 => GradTensor::D5(unsafe { std::mem::transmute_copy(&grad_input) }),
+            5 => GradTensor::D5(into_dim::<B, D, 5>(grad_input)),
-            6 => GradTensor::D6(unsafe { std::mem::transmute_copy(&grad_input) }),
+            6 => GradTensor::D6(into_dim::<B, D, 6>(grad_input)),
             _ => return Err(AutogradError::UnsupportedDimension(0)),
         };
 
@@ -940,12 +945,12 @@ where
 
         // SAFETY: See module-level documentation. D is verified by each match arm.
         let result = match D {
-            1 => GradTensor::D1(unsafe { std::mem::transmute_copy(&grad_input) }),
+            1 => GradTensor::D1(into_dim::<B, D, 1>(grad_input)),
-            2 => GradTensor::D2(unsafe { std::mem::transmute_copy(&grad_input) }),
+            2 => GradTensor::D2(into_dim::<B, D, 2>(grad_input)),
-            3 => GradTensor::D3(unsafe { std::mem::transmute_copy(&grad_input) }),
+            3 => GradTensor::D3(into_dim::<B, D, 3>(grad_input)),
-            4 => GradTensor::D4(unsafe { std::mem::transmute_copy(&grad_input) }),
+            4 => GradTensor::D4(into_dim::<B, D, 4>(grad_input)),
-            5 => GradTensor::D5(unsafe { std::mem::transmute_copy(&grad_input) }),
+            5 => GradTensor::D5(into_dim::<B, D, 5>(grad_input)),
-            6 => GradTensor::D6(unsafe { std::mem::transmute_copy(&grad_input) }),
+            6 => GradTensor::D6(into_dim::<B, D, 6>(grad_input)),
             _ => return Err(AutogradError::UnsupportedDimension(0)),
         };
 

crates/core/rtx-autograd/src/autodiff/ops/shapes.rs

@@ -7,6 +7,7 @@
 //! - SwapDims: gradient swaps the same dimensions back
 
 use crate::autodiff::node::{AutodiffBackwardFn, GradTensor, SavedTensor};
+use crate::autodiff::ops::into_dim;
 use crate::error::{AutogradError, Result};
 use rtx_backend::Backend;
 use std::marker::PhantomData;
@@ -214,17 +215,17 @@ fn wrap_as_grad_tensor<B: Backend, const D: usize>(
 ) -> Result<GradTensor<B>> {
     match D {
         // SAFETY: D=1 verified by match, B::TensorPrimitive layout identical across D
-        1 => Ok(GradTensor::D1(unsafe { std::mem::transmute_copy(&tensor) })),
+        1 => Ok(GradTensor::D1(into_dim::<B, D, 1>(tensor))),
         // SAFETY: D=2 verified by match, B::TensorPrimitive layout identical across D
-        2 => Ok(GradTensor::D2(unsafe { std::mem::transmute_copy(&tensor) })),
+        2 => Ok(GradTensor::D2(into_dim::<B, D, 2>(tensor))),
         // SAFETY: D=3 verified by match, B::TensorPrimitive layout identical across D
-        3 => Ok(GradTensor::D3(unsafe { std::mem::transmute_copy(&tensor) })),
+        3 => Ok(GradTensor::D3(into_dim::<B, D, 3>(tensor))),
         // SAFETY: D=4 verified by match, B::TensorPrimitive layout identical across D
-        4 => Ok(GradTensor::D4(unsafe { std::mem::transmute_copy(&tensor) })),
+        4 => Ok(GradTensor::D4(into_dim::<B, D, 4>(tensor))),
         // SAFETY: D=5 verified by match, B::TensorPrimitive layout identical across D
-        5 => Ok(GradTensor::D5(unsafe { std::mem::transmute_copy(&tensor) })),
+        5 => Ok(GradTensor::D5(into_dim::<B, D, 5>(tensor))),
         // SAFETY: D=6 verified by match, B::TensorPrimitive layout identical across D
-        6 => Ok(GradTensor::D6(unsafe { std::mem::transmute_copy(&tensor) })),
+        6 => Ok(GradTensor::D6(into_dim::<B, D, 6>(tensor))),
         _ => Err(AutogradError::UnsupportedDimension(D)),
     }
 }

crates/core/rtx-autograd/src/autodiff/ops/unary.rs

@@ -8,6 +8,7 @@
 //! - Abs: `grad_input = grad_output * sign(input)`
 
 use crate::autodiff::node::{AutodiffBackwardFn, GradTensor, SavedTensor};
+use crate::autodiff::ops::into_dim;
 use crate::error::{AutogradError, Result};
 use rtx_backend::Backend;
 use std::marker::PhantomData;
@@ -255,32 +256,32 @@ where
     match (grad, D) {
         (GradTensor::D1(g), 1) => {
             // SAFETY: D=1 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<1>>(tensor) };
+            let t = into_dim::<B, D, 1>(tensor.clone());
             Ok(GradTensor::D1(B::mul(g, t)))
         }
         (GradTensor::D2(g), 2) => {
             // SAFETY: D=2 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<2>>(tensor) };
+            let t = into_dim::<B, D, 2>(tensor.clone());
             Ok(GradTensor::D2(B::mul(g, t)))
         }
         (GradTensor::D3(g), 3) => {
             // SAFETY: D=3 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<3>>(tensor) };
+            let t = into_dim::<B, D, 3>(tensor.clone());
             Ok(GradTensor::D3(B::mul(g, t)))
         }
         (GradTensor::D4(g), 4) => {
             // SAFETY: D=4 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<4>>(tensor) };
+            let t = into_dim::<B, D, 4>(tensor.clone());
             Ok(GradTensor::D4(B::mul(g, t)))
         }
         (GradTensor::D5(g), 5) => {
             // SAFETY: D=5 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<5>>(tensor) };
+            let t = into_dim::<B, D, 5>(tensor.clone());
             Ok(GradTensor::D5(B::mul(g, t)))
         }
         (GradTensor::D6(g), 6) => {
             // SAFETY: D=6 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<6>>(tensor) };
+            let t = into_dim::<B, D, 6>(tensor.clone());
             Ok(GradTensor::D6(B::mul(g, t)))
         }
         _ => Err(AutogradError::DimensionMismatch("mul_grad".to_string())),
@@ -304,32 +305,32 @@ where
     match (grad, D) {
         (GradTensor::D1(g), 1) => {
             // SAFETY: D=1 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<1>>(tensor) };
+            let t = into_dim::<B, D, 1>(tensor.clone());
             Ok(GradTensor::D1(B::div(g, t)))
         }
         (GradTensor::D2(g), 2) => {
             // SAFETY: D=2 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<2>>(tensor) };
+            let t = into_dim::<B, D, 2>(tensor.clone());
             Ok(GradTensor::D2(B::div(g, t)))
         }
         (GradTensor::D3(g), 3) => {
             // SAFETY: D=3 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<3>>(tensor) };
+            let t = into_dim::<B, D, 3>(tensor.clone());
             Ok(GradTensor::D3(B::div(g, t)))
         }
         (GradTensor::D4(g), 4) => {
             // SAFETY: D=4 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<4>>(tensor) };
+            let t = into_dim::<B, D, 4>(tensor.clone());
             Ok(GradTensor::D4(B::div(g, t)))
         }
         (GradTensor::D5(g), 5) => {
             // SAFETY: D=5 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<5>>(tensor) };
+            let t = into_dim::<B, D, 5>(tensor.clone());
             Ok(GradTensor::D5(B::div(g, t)))
         }
         (GradTensor::D6(g), 6) => {
             // SAFETY: D=6 verified by match, layout identical across D values
-            let t = unsafe { std::mem::transmute_copy::<_, B::TensorPrimitive<6>>(tensor) };
+            let t = into_dim::<B, D, 6>(tensor.clone());
             Ok(GradTensor::D6(B::div(g, t)))
         }
         _ => Err(AutogradError::DimensionMismatch("div_grad".to_string())),

crates/core/rtx-autograd/src/autodiff/tensor.rs

@@ -169,7 +169,16 @@ where
         Self {
             inner: self.inner.clone(),
             node: self.node.clone(), // Arc clone is cheap
-            id: TensorId::new(),     // New ID for the clone
+            // Preserve the id: in this graph model the id identifies a node
+            // in the autograd graph, so a clone is the *same* logical tensor
+            // and must share the same gradient sink. Minting a fresh id here
+            // meant that using a tensor twice (which requires a clone, since
+            // ops take their operands by value) split its gradient across two
+            // ids; backward then accumulated into both separately and the
+            // caller, querying only the original id, saw a fraction of the
+            // true gradient. Sharing the id lets fan-out paths collide on one
+            // id and sum via `accumulate_gradients`.
+            id: self.id,
             requires_grad: self.requires_grad,
         }
     }

crates/core/rtx-autograd/tests/tape_cpu_gradcheck.rs

@@ -0,0 +1,174 @@
+//! Numerical finite-difference gradient checks for the `Autodiff<B>` tape,
+//! run against the **real** `CpuBackend` — not the shape-only `MockBackend`
+//! the rest of the suite uses (whose ops all return their input, so they
+//! validate graph *structure* but never gradient *values*).
+//!
+//! This is the first end-to-end validation that the decorator tape produces
+//! correct numerical gradients through a concrete backend, which is the
+//! precondition for training a tape-native transformer (the SMT teacher).
+//!
+//! Method: build a scalar loss through the tape, extract the leaf gradient
+//! via `backward_impl` + `GradientStorage`, and compare every element to a
+//! central finite difference computed with the raw (grad-free) `CpuBackend`.
+//!
+//! Note on Miri: these pass under normal `cargo test`. Under `cargo miri test`
+//! they abort on a Stacked-Borrows / integer-to-pointer violation that lives
+//! in `rtx-backend-cpu`'s buffer internals, NOT in the autodiff tape — a
+//! grad-free `from_data`+`add`+`sum` probe reproduces the identical error, so
+//! the autograd path is exonerated. Fixing that backend UB is tracked
+//! separately.
+
+use rtx_autograd::autodiff::{Autodiff, AutodiffDevice, GradTensor, backward_impl};
+use rtx_backend::{AutodiffBackend, Backend};
+use rtx_backend_cpu::CpuBackend;
+
+type Ad = Autodiff<CpuBackend>;
+
+const H: f32 = 1e-3;
+const TOL: f32 = 2e-2;
+
+fn ad_dev() -> AutodiffDevice<CpuBackend> {
+    AutodiffDevice::<CpuBackend>::default()
+}
+
+fn cpu_dev() -> <CpuBackend as Backend>::Device {
+    <CpuBackend as Backend>::Device::default()
+}
+
+/// Make a grad-tracked 2-D leaf on the tape.
+fn leaf2(data: &[f32], shape: [usize; 2]) -> <Ad as Backend>::TensorPrimitive<2> {
+    Ad::require_grad(Ad::from_data(data, shape, &ad_dev()))
+}
+
+/// Extract a leaf's 2-D gradient from the storage as a flat Vec.
+fn grad2(storage: &rtx_autograd::autodiff::GradientStorage<CpuBackend>, id: usize) -> Vec<f32> {
+    let g = storage
+        .get(rtx_autograd::autodiff::TensorId(id))
+        .expect("gradient present for leaf");
+    match g {
+        GradTensor::D2(t) => t.to_vec(),
+        other => panic!("expected D2 gradient, got {other:?}"),
+    }
+}
+
+/// Central finite-difference gradient of a scalar loss `f(a)` wrt each
+/// element of `a`, computed with the raw grad-free backend.
+fn fd_grad(a: &[f32], f: impl Fn(&[f32]) -> f32) -> Vec<f32> {
+    let mut g = vec![0.0f32; a.len()];
+    let mut probe = a.to_vec();
+    for i in 0..a.len() {
+        let orig = probe[i];
+        probe[i] = orig + H;
+        let plus = f(&probe);
+        probe[i] = orig - H;
+        let minus = f(&probe);
+        probe[i] = orig;
+        g[i] = (plus - minus) / (2.0 * H);
+    }
+    g
+}
+
+fn assert_close(analytical: &[f32], numerical: &[f32], what: &str) {
+    assert_eq!(analytical.len(), numerical.len(), "{what}: length mismatch");
+    for (i, (a, n)) in analytical.iter().zip(numerical.iter()).enumerate() {
+        let denom = a.abs().max(n.abs()).max(1.0);
+        let rel = (a - n).abs() / denom;
+        assert!(
+            rel < TOL,
+            "{what}: grad[{i}] analytical={a} numerical={n} rel_err={rel}"
+        );
+    }
+}
+
+/// `loss = sum(A @ B)` — checks the matmul VJP for both operands plus the
+/// sum VJP and the grad-extraction path.
+#[test]
+fn gradcheck_matmul_sum() {
+    let a = [0.5f32, -1.0, 2.0, 0.25, -0.5, 1.5]; // [2,3]
+    let b = [1.0f32, -2.0, 0.5, 3.0, -1.5, 0.75]; // [3,2]
+
+    // Raw forward used by finite differences (vary A, fix B).
+    let loss_of_a = |av: &[f32]| -> f32 {
+        let at = CpuBackend::from_data(av, [2, 3], &cpu_dev());
+        let bt = CpuBackend::from_data(&b, [3, 2], &cpu_dev());
+        CpuBackend::to_data(&CpuBackend::sum(CpuBackend::matmul(at, bt)))[0]
+    };
+
+    let at = leaf2(&a, [2, 3]);
+    let bt = leaf2(&b, [3, 2]);
+    let a_id = at.id().0;
+    let loss = Ad::sum(Ad::matmul(at, bt));
+    let storage = backward_impl(
+        &loss,
+        Some(GradTensor::from_d1(CpuBackend::ones([1], &cpu_dev()))),
+    )
+    .expect("backward");
+
+    let analytical = grad2(&storage, a_id);
+    let numerical = fd_grad(&a, loss_of_a);
+    assert_close(&analytical, &numerical, "matmul_sum d/dA");
+}
+
+/// `loss = sum((A + B) * (A + B))` — checks same-shape add VJP and elementwise
+/// mul VJP, and specifically the **fan-out accumulation** path: `s` feeds both
+/// operands of `mul(s.clone(), s)`, so its gradient arrives along two paths and
+/// must be summed. This previously came out exactly halved because
+/// `accumulate_gradients` discarded one path and `clone` minted a fresh id;
+/// fixed by summing via `B::add` and preserving the id across `clone`.
+#[test]
+fn gradcheck_add_mul_sum() {
+    let a = [0.5f32, -1.0, 2.0, 0.25]; // [2,2]
+    let b = [1.0f32, -2.0, 0.5, 3.0]; // [2,2]
+
+    let loss_of_a = |av: &[f32]| -> f32 {
+        let at = CpuBackend::from_data(av, [2, 2], &cpu_dev());
+        let bt = CpuBackend::from_data(&b, [2, 2], &cpu_dev());
+        let s = CpuBackend::add(at, bt);
+        CpuBackend::to_data(&CpuBackend::sum(CpuBackend::mul(s.clone(), s)))[0]
+    };
+
+    let at = leaf2(&a, [2, 2]);
+    let bt = leaf2(&b, [2, 2]);
+    let a_id = at.id().0;
+    let s = Ad::add(at, bt);
+    let loss = Ad::sum(Ad::mul(s.clone(), s));
+    let storage = backward_impl(
+        &loss,
+        Some(GradTensor::from_d1(CpuBackend::ones([1], &cpu_dev()))),
+    )
+    .expect("backward");
+
+    let analytical = grad2(&storage, a_id);
+    let numerical = fd_grad(&a, loss_of_a);
+    assert_close(&analytical, &numerical, "add_mul_sum d/dA");
+}
+
+/// `loss = sum(softmax(X, dim=1) ^ 2)` — the riskiest transformer VJP. This
+/// previously panicked in the CPU `sub` because softmax backward subtracted a
+/// keep-dim `[.., 1]` row-sum from the full `[.., cols]` grad, and the
+/// elementwise ops require equal shapes; fixed by tiling the row-sum to full
+/// width (`broadcast_along_dim`) before the subtract.
+#[test]
+fn gradcheck_softmax_sum() {
+    let x = [0.2f32, -0.5, 1.0, 0.3, 0.8, -1.2]; // [2,3]
+
+    let loss_of_x = |xv: &[f32]| -> f32 {
+        let xt = CpuBackend::from_data(xv, [2, 3], &cpu_dev());
+        let y = CpuBackend::softmax(xt, 1);
+        CpuBackend::to_data(&CpuBackend::sum(CpuBackend::mul(y.clone(), y)))[0]
+    };
+
+    let xt = leaf2(&x, [2, 3]);
+    let x_id = xt.id().0;
+    let y = Ad::softmax(xt, 1);
+    let loss = Ad::sum(Ad::mul(y.clone(), y));
+    let storage = backward_impl(
+        &loss,
+        Some(GradTensor::from_d1(CpuBackend::ones([1], &cpu_dev()))),
+    )
+    .expect("backward");
+
+    let analytical = grad2(&storage, x_id);
+    let numerical = fd_grad(&x, loss_of_x);
+    assert_close(&analytical, &numerical, "softmax_sum d/dX");
+}