feat(backend): add differentiable index_select / index_add row ops

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

crates/core/rtx-backend/src/lib.rs

@@ -280,6 +280,84 @@ pub trait Backend: Clone + Send + Sync + Debug + Default + 'static {
         dim2: usize,
     ) -> Self::TensorPrimitive<D>;
 
+    // ==================== Row Indexing (gather / scatter-add) ====================
+    // Differentiable row indexing along dim 0. These two ops are each other's
+    // adjoint, which is exactly what message passing on a graph needs:
+    //   d/dx index_select(x, idx)        = index_add(grad, idx, rows(x))
+    //   d/dx index_add(x, idx, num_rows) = index_select(grad, idx)
+    //
+    // Both have default bodies that round-trip through host memory via
+    // `to_data` / `from_data`, so every backend is correct out of the box;
+    // backends override them with native kernels for speed.
+
+    /// Gather rows along dim 0: `out[i, ..] = tensor[indices[i], ..]`.
+    ///
+    /// Output shape is `[indices.len(), shape[1..]]`. Indices may repeat.
+    ///
+    /// # Panics
+    /// Panics if any index is `>= shape[0]`, or if `D == 0`.
+    fn index_select<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        indices: &[usize],
+    ) -> Self::TensorPrimitive<D> {
+        let shape = Self::shape(&tensor);
+        assert!(D >= 1, "index_select requires at least one dimension");
+        let num_rows = shape[0];
+        let row_len: usize = shape[1..].iter().product();
+        let src = Self::to_data(&tensor);
+        let mut out = Vec::with_capacity(indices.len() * row_len);
+        for &idx in indices {
+            assert!(
+                idx < num_rows,
+                "index_select: index {idx} out of range for {num_rows} rows"
+            );
+            out.extend_from_slice(&src[idx * row_len..(idx + 1) * row_len]);
+        }
+        let mut out_shape = shape;
+        out_shape[0] = indices.len();
+        Self::from_data(&out, out_shape, &Self::device(&tensor))
+    }
+
+    /// Scatter-add rows along dim 0 into a zero tensor with `num_rows` rows:
+    /// `out = zeros([num_rows, shape[1..]]); out[indices[i], ..] += tensor[i, ..]`.
+    ///
+    /// Indices may repeat (contributions accumulate); rows never referenced
+    /// stay zero. This is the adjoint of [`Backend::index_select`].
+    ///
+    /// # Panics
+    /// Panics if `indices.len() != shape[0]`, if any index is `>= num_rows`,
+    /// or if `D == 0`.
+    fn index_add<const D: usize>(
+        tensor: Self::TensorPrimitive<D>,
+        indices: &[usize],
+        num_rows: usize,
+    ) -> Self::TensorPrimitive<D> {
+        let shape = Self::shape(&tensor);
+        assert!(D >= 1, "index_add requires at least one dimension");
+        assert_eq!(
+            indices.len(),
+            shape[0],
+            "index_add: indices.len() must equal the number of input rows"
+        );
+        let row_len: usize = shape[1..].iter().product();
+        let src = Self::to_data(&tensor);
+        let mut out = vec![Self::FloatElem::zero(); num_rows * row_len];
+        for (i, &idx) in indices.iter().enumerate() {
+            assert!(
+                idx < num_rows,
+                "index_add: index {idx} out of range for {num_rows} rows"
+            );
+            let dst = &mut out[idx * row_len..(idx + 1) * row_len];
+            let row = &src[i * row_len..(i + 1) * row_len];
+            for (d, &s) in dst.iter_mut().zip(row) {
+                *d = Self::FloatElem::from_f64(d.to_f64() + s.to_f64());
+            }
+        }
+        let mut out_shape = shape;
+        out_shape[0] = num_rows;
+        Self::from_data(&out, out_shape, &Self::device(&tensor))
+    }
+
     // ==================== LLM-Specific Operations ====================
     // These delegate to hand-optimized kernels for maximum performance.