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]>
Add two row-indexing ops along dim 0 to the `Backend` trait so gather /
scatter-add message passing (GNNs, segment softmax, bias tiling) can be
trained through `Autodiff<B>`:
- `index_select(tensor, indices)` — out[i, ..] = tensor[indices[i], ..]
- `index_add(tensor, indices, num_rows)` — out = zeros; out[idx[i], ..] += tensor[i, ..]
They are each other's adjoint, which is what the backward passes use.
Both trait methods have default bodies (host round-trip via to_data /
from_data) so every existing backend keeps compiling and is correct;
backends override with native kernels:
- rtx-backend-cpu: new ops/index.rs (rayon-parallel gather over output
rows above a size threshold, sequential deterministic scatter-add),
wired into CpuBackend and CpuBackendF64, with unit tests for D=1/2/3,
duplicates, untouched rows, empty inputs, bounds panics and adjointness.
- rtx-autograd: Autodiff<B> overrides both ops and records
IndexSelectBackward / IndexAddBackward (new ops/index.rs); finite-
difference gradchecks on the real CpuBackend cover repeated-index
accumulation, untouched-row zero grads, bias tiling via index_select
of a [1,F] row, and a full per-segment softmax.
- rtx-fusion: forward both ops to the inner backend.
Co-Authored-By: Claude Fable 5 <[email protected]>
The decorator autograd (`Autodiff<B>`) had never been gradient-checked
against a real tensor backend — the entire test suite runs on a shape-only
`MockBackend` whose ops return their input, so they validate graph structure
but never gradient values. Running it through `CpuBackend` for the first time
(new `tests/tape_cpu_gradcheck.rs`, finite-difference checks) surfaced three
bugs that made the tape unusable for training; this fixes all three.
1. Double-free / UB in the dimension-erasure cast. The backward ops cast a
tensor to its runtime const-generic dimension via
`mem::transmute_copy::<_, TensorPrimitive<N>>(&src)` in ~100 sites. That
bit-copies the owned `Vec` without forgetting the source, so two values own
one buffer → double-free on any heap-backed backend (and Stacked-Borrows UB
from the typed pun). Replaced every site with a single `into_dim` helper
that is now **fully safe** — it round-trips through `to_data`/`from_data`
and rebuilds the shape with `array::from_fn`, no `unsafe` at all. (This is
why the whole repo previously bypassed the tape with analytic backward.)
2. Fan-out gradients were silently dropped. `accumulate_gradients` was a stub
that returned one path and discarded the other, and `AutodiffTensor::clone`
minted a fresh `TensorId`. Together, reusing a tensor (residuals,
`mul(s, s)`, shared Q/K/V — universal in transformers) split its gradient
across two ids and summed neither, yielding a fraction of the true value.
`accumulate_gradients` now sums via `B::add`; `clone` preserves the id so
fan-out paths collide on one sink.
3. Softmax backward panicked. `SoftmaxBackward` / `stable_softmax_backward`
subtracted a keep-dim row-sum from the full-shape grad, but the elementwise
backends assert equal shapes (no broadcasting). Added `broadcast_along_dim`
to tile the row-sum to full width first.
Verified: `tape_cpu_gradcheck` (matmul, fan-out add·mul, softmax) passes with
rel-err < 2e-2 vs central differences; full `rtx-autograd` suite green (263
passed, 0 failed); lib clippy `-D warnings` clean.
Known follow-up (out of scope): `cargo miri test` still aborts on a
Stacked-Borrows / integer-to-pointer violation inside `rtx-backend-cpu`'s
buffer internals — a grad-free `from_data`+`add`+`sum` probe reproduces the
identical error, so it is pre-existing backend UB, not an autograd issue.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
osobhandClaude Opus 4.8