fix(autograd): AbsBackward produced NaN for exactly-zero inputs

sign was computed as x/|x|, which is 0/0 = NaN at x = 0; one zero element
in an |pred - target| loss poisoned every upstream gradient (hit
deterministically by dg-gnn HetGAT training). Compute x/(|x| + tiny)
instead so sign(0) = 0 (the subgradient convention). Regression test
included.

Co-Authored-By: Claude Fable 5 <[email protected]>

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

@@ -218,11 +218,20 @@ where
             .downcast_ref::<B::TensorPrimitive<D>>()
             .ok_or_else(|| AutogradError::DowncastError("saved input".to_string()))?;
 
-        // Compute sign: x / |x|
+        // Compute sign as x / (|x| + tiny): -1 for negative, +1 for positive,
-        // This gives -1 for negative, +1 for positive, and 0/0=NaN for zero
+        // and — crucially — 0 for zero. The previous x / |x| form produced
-        // We handle this by computing: grad * (input / abs(input))
+        // 0/0 = NaN whenever an input element was exactly zero (|x| is
+        // non-differentiable there; the subgradient convention is sign(0)=0),
+        // and a single such element poisoned every upstream gradient. The
+        // tiny denominator offset only perturbs the sign of subnormal inputs.
+        use rtx_backend::FloatElement;
         let abs_input = B::abs(input.clone());
-        let sign = B::div(input.clone(), abs_input);
+        let tiny = B::full(
+            B::shape(input),
+            B::FloatElem::from_f64(f64::from(f32::MIN_POSITIVE)),
+            &B::device(input),
+        );
+        let sign = B::div(input.clone(), B::add(abs_input, tiny));
 
         // grad_input = grad_output * sign
         let grad_input = mul_grad::<B, D>(grad_output, &sign)?;

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

@@ -496,3 +496,28 @@ fn gradcheck_segment_softmax() {
     let numerical = fd_grad(&x, loss_of_x);
     assert_close(&analytical, &numerical, "segment_softmax d/dX");
 }
+
+/// `abs` backward at an exactly-zero input element must yield gradient 0 for
+/// that element (subgradient convention), never NaN. Regression: the previous
+/// sign = x/|x| produced 0/0 = NaN and poisoned every upstream gradient
+/// (surfaced by dg-gnn HetGAT training, where |pred − target| hits exact
+/// zeros over long runs).
+#[test]
+fn abs_backward_zero_input_is_zero_not_nan() {
+    let x = [0.5f32, 0.0, -2.0, 0.0, 1.0e-30, -0.0];
+    let xt = leaf2(&x, [3, 2]);
+    let x_id = xt.id().0;
+    let loss = Ad::sum(Ad::abs(xt));
+    let storage = backward_impl(
+        &loss,
+        Some(GradTensor::from_d1(CpuBackend::ones([1], &cpu_dev()))),
+    )
+    .expect("backward");
+    let g = grad2(&storage, x_id);
+    assert!(g.iter().all(|v| v.is_finite()), "abs grad has non-finite values: {g:?}");
+    assert_eq!(g[0], 1.0);
+    assert_eq!(g[1], 0.0, "sign(0) must be 0");
+    assert_eq!(g[2], -1.0);
+    assert_eq!(g[3], 0.0);
+    assert_eq!(g[5], 0.0, "sign(-0.0) must be 0");
+}