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]>
