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]>
This commit is contained in:
Omar Sobh
co-authored by
Claude Fable 5
Claude Fable 5
parent
327da7ff47
commit
9969d8a661
@@ -0,0 +1,265 @@
|
||||
//! Row indexing operations: gather (`index_select`) and scatter-add (`index_add`).
|
||||
//!
|
||||
//! Both operate along dim 0 and treat every trailing dimension as a flat,
|
||||
//! contiguous "row" of `shape[1..].product()` elements. They are each other's
|
||||
//! adjoint, which is what lets a graph message-passing layer be trained with
|
||||
//! `Autodiff<CpuBackend>`:
|
||||
//!
|
||||
//! - `index_select(x, idx)[i, ..] = x[idx[i], ..]`
|
||||
//! - `index_add(x, idx, n)[idx[i], ..] += x[i, ..]` (starting from zeros)
|
||||
//!
|
||||
//! `CpuTensorPrimitive` data is always stored contiguously in row-major order
|
||||
//! (every op materialises a fresh contiguous buffer), so slicing rows directly
|
||||
//! out of `data` is valid.
|
||||
|
||||
use crate::{CpuFloat, CpuTensorPrimitive};
|
||||
use rayon::prelude::*;
|
||||
|
||||
/// Minimum number of output elements before the gather goes parallel.
|
||||
const PAR_THRESHOLD: usize = 1 << 12;
|
||||
|
||||
/// 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 out of range or `D == 0`.
|
||||
pub fn index_select<const D: usize, E: CpuFloat>(
|
||||
tensor: &CpuTensorPrimitive<D, E>,
|
||||
indices: &[usize],
|
||||
) -> CpuTensorPrimitive<D, E> {
|
||||
assert!(D >= 1, "index_select requires at least one dimension");
|
||||
let num_rows = tensor.shape[0];
|
||||
let row_len: usize = tensor.shape[1..].iter().product();
|
||||
for &idx in indices {
|
||||
assert!(
|
||||
idx < num_rows,
|
||||
"index_select: index {idx} out of range for {num_rows} rows"
|
||||
);
|
||||
}
|
||||
|
||||
let src = &tensor.data;
|
||||
let total = indices.len() * row_len;
|
||||
let mut out: Vec<E> = vec![E::zero(); total];
|
||||
|
||||
if row_len > 0 {
|
||||
if total >= PAR_THRESHOLD {
|
||||
out.par_chunks_mut(row_len)
|
||||
.zip(indices.par_iter())
|
||||
.for_each(|(dst, &idx)| {
|
||||
dst.copy_from_slice(&src[idx * row_len..(idx + 1) * row_len]);
|
||||
});
|
||||
} else {
|
||||
for (dst, &idx) in out.chunks_mut(row_len).zip(indices) {
|
||||
dst.copy_from_slice(&src[idx * row_len..(idx + 1) * row_len]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut out_shape = tensor.shape;
|
||||
out_shape[0] = indices.len();
|
||||
CpuTensorPrimitive::new(out, out_shape, tensor.device.clone())
|
||||
}
|
||||
|
||||
/// Scatter-add rows along dim 0 into a zero tensor with `num_rows` rows:
|
||||
/// `out[indices[i], ..] += tensor[i, ..]`.
|
||||
///
|
||||
/// Indices may repeat (contributions accumulate); rows never referenced stay
|
||||
/// zero. Sequential so that repeated indices accumulate deterministically.
|
||||
///
|
||||
/// # Panics
|
||||
/// Panics if `indices.len() != shape[0]`, any index is `>= num_rows`, or `D == 0`.
|
||||
pub fn index_add<const D: usize, E: CpuFloat>(
|
||||
tensor: &CpuTensorPrimitive<D, E>,
|
||||
indices: &[usize],
|
||||
num_rows: usize,
|
||||
) -> CpuTensorPrimitive<D, E> {
|
||||
assert!(D >= 1, "index_add requires at least one dimension");
|
||||
assert_eq!(
|
||||
indices.len(),
|
||||
tensor.shape[0],
|
||||
"index_add: indices.len() must equal the number of input rows"
|
||||
);
|
||||
let row_len: usize = tensor.shape[1..].iter().product();
|
||||
let src = &tensor.data;
|
||||
let mut out: Vec<E> = vec![E::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 = *d + s;
|
||||
}
|
||||
}
|
||||
|
||||
let mut out_shape = tensor.shape;
|
||||
out_shape[0] = num_rows;
|
||||
CpuTensorPrimitive::new(out, out_shape, tensor.device.clone())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::CpuDevice;
|
||||
|
||||
fn t2(data: &[f32], shape: [usize; 2]) -> CpuTensorPrimitive<2> {
|
||||
CpuTensorPrimitive::new(data.to_vec(), shape, CpuDevice::new())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn index_select_2d_gathers_rows_with_duplicates() {
|
||||
// 3 rows x 2 cols
|
||||
let x = t2(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [3, 2]);
|
||||
let y = index_select(&x, &[2, 0, 2, 1]);
|
||||
assert_eq!(y.shape(), &[4, 2]);
|
||||
assert_eq!(y.to_vec(), vec![5.0, 6.0, 1.0, 2.0, 5.0, 6.0, 3.0, 4.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn index_select_empty_indices_gives_zero_rows() {
|
||||
let x = t2(&[1.0, 2.0, 3.0, 4.0], [2, 2]);
|
||||
let y = index_select(&x, &[]);
|
||||
assert_eq!(y.shape(), &[0, 2]);
|
||||
assert!(y.to_vec().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn index_select_1d_gathers_scalars() {
|
||||
let x = CpuTensorPrimitive::<1>::new(vec![10.0, 20.0, 30.0], [3], CpuDevice::new());
|
||||
let y = index_select(&x, &[1, 1, 0]);
|
||||
assert_eq!(y.shape(), &[3]);
|
||||
assert_eq!(y.to_vec(), vec![20.0, 20.0, 10.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn index_select_3d_gathers_whole_slabs() {
|
||||
// shape [2, 2, 2]
|
||||
let x = CpuTensorPrimitive::<3>::new(
|
||||
(0..8).map(|v| v as f32).collect(),
|
||||
[2, 2, 2],
|
||||
CpuDevice::new(),
|
||||
);
|
||||
let y = index_select(&x, &[1, 0]);
|
||||
assert_eq!(y.shape(), &[2, 2, 2]);
|
||||
assert_eq!(y.to_vec(), vec![4.0, 5.0, 6.0, 7.0, 0.0, 1.0, 2.0, 3.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn index_select_large_goes_parallel_and_matches() {
|
||||
let rows = 300;
|
||||
let cols = 32;
|
||||
let data: Vec<f32> = (0..rows * cols).map(|v| v as f32).collect();
|
||||
let x = t2(&data, [rows, cols]);
|
||||
let idx: Vec<usize> = (0..rows * 2).map(|i| (i * 7) % rows).collect();
|
||||
let y = index_select(&x, &idx);
|
||||
assert_eq!(y.shape(), &[rows * 2, cols]);
|
||||
let out = y.to_vec();
|
||||
for (i, &r) in idx.iter().enumerate() {
|
||||
assert_eq!(
|
||||
&out[i * cols..(i + 1) * cols],
|
||||
&data[r * cols..(r + 1) * cols]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "out of range")]
|
||||
fn index_select_out_of_range_panics() {
|
||||
let x = t2(&[1.0, 2.0, 3.0, 4.0], [2, 2]);
|
||||
let _ = index_select(&x, &[2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn index_add_2d_accumulates_duplicates_and_leaves_untouched_zero() {
|
||||
// 4 input rows x 2 cols scattered into 3 output rows; row 1 untouched.
|
||||
let x = t2(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], [4, 2]);
|
||||
let y = index_add(&x, &[2, 0, 2, 0], 3);
|
||||
assert_eq!(y.shape(), &[3, 2]);
|
||||
assert_eq!(
|
||||
y.to_vec(),
|
||||
vec![3.0 + 7.0, 4.0 + 8.0, 0.0, 0.0, 1.0 + 5.0, 2.0 + 6.0]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn index_add_1d_accumulates() {
|
||||
let x = CpuTensorPrimitive::<1>::new(vec![1.0, 2.0, 3.0], [3], CpuDevice::new());
|
||||
let y = index_add(&x, &[1, 1, 0], 4);
|
||||
assert_eq!(y.shape(), &[4]);
|
||||
assert_eq!(y.to_vec(), vec![3.0, 3.0, 0.0, 0.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn index_add_3d_accumulates_slabs() {
|
||||
let x = CpuTensorPrimitive::<3>::new(
|
||||
(0..8).map(|v| v as f32).collect(),
|
||||
[2, 2, 2],
|
||||
CpuDevice::new(),
|
||||
);
|
||||
let y = index_add(&x, &[0, 0], 2);
|
||||
assert_eq!(y.shape(), &[2, 2, 2]);
|
||||
assert_eq!(y.to_vec(), vec![4.0, 6.0, 8.0, 10.0, 0.0, 0.0, 0.0, 0.0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn index_add_empty_input_gives_zeros() {
|
||||
let x = t2(&[], [0, 3]);
|
||||
let y = index_add(&x, &[], 2);
|
||||
assert_eq!(y.shape(), &[2, 3]);
|
||||
assert_eq!(y.to_vec(), vec![0.0; 6]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "indices.len()")]
|
||||
fn index_add_len_mismatch_panics() {
|
||||
let x = t2(&[1.0, 2.0, 3.0, 4.0], [2, 2]);
|
||||
let _ = index_add(&x, &[0], 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "out of range")]
|
||||
fn index_add_out_of_range_panics() {
|
||||
let x = t2(&[1.0, 2.0, 3.0, 4.0], [2, 2]);
|
||||
let _ = index_add(&x, &[0, 5], 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn select_and_add_are_adjoint() {
|
||||
// <index_select(x, idx), g> == <x, index_add(g, idx, n)>
|
||||
let n = 5;
|
||||
let cols = 3;
|
||||
let x_data: Vec<f32> = (0..n * cols).map(|v| (v as f32) * 0.5 - 2.0).collect();
|
||||
let x = t2(&x_data, [n, cols]);
|
||||
let idx = [4usize, 0, 4, 2, 2, 1];
|
||||
let g_data: Vec<f32> = (0..idx.len() * cols).map(|v| (v as f32).sin()).collect();
|
||||
let g = t2(&g_data, [idx.len(), cols]);
|
||||
|
||||
let lhs: f32 = index_select(&x, &idx)
|
||||
.to_vec()
|
||||
.iter()
|
||||
.zip(&g_data)
|
||||
.map(|(a, b)| a * b)
|
||||
.sum();
|
||||
let rhs: f32 = index_add(&g, &idx, n)
|
||||
.to_vec()
|
||||
.iter()
|
||||
.zip(&x_data)
|
||||
.map(|(a, b)| a * b)
|
||||
.sum();
|
||||
assert!((lhs - rhs).abs() < 1e-4, "lhs={lhs} rhs={rhs}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn f64_path_works() {
|
||||
let x =
|
||||
CpuTensorPrimitive::<2, f64>::new(vec![1.0, 2.0, 3.0, 4.0], [2, 2], CpuDevice::new());
|
||||
let y = index_select(&x, &[1, 1]);
|
||||
assert_eq!(y.to_vec(), vec![3.0, 4.0, 3.0, 4.0]);
|
||||
let z = index_add(&y, &[0, 0], 3);
|
||||
assert_eq!(z.to_vec(), vec![6.0, 8.0, 0.0, 0.0, 0.0, 0.0]);
|
||||
}
|
||||
}
|
||||
@@ -8,6 +8,7 @@ pub mod basic;
|
||||
pub mod conv;
|
||||
pub mod creation;
|
||||
pub mod gemm;
|
||||
pub mod index;
|
||||
pub mod normalization;
|
||||
pub mod pooling;
|
||||
pub mod reduction;
|
||||
|
||||
Reference in New Issue
Block a user