nuvec_core/
indexing.rs

1//! Indexing and joining: `take`, `take_along_axis`, `put_along_axis`, `concatenate`,
2//! `stack` and `split`, as in MLX.
3//!
4//! Indices are integer arrays. As in MLX, negative indices count from the end of the
5//! axis; out-of-range indices (which MLX leaves unspecified) are clamped to the axis.
6
7use crate::ops::{broadcast_to, expand_dims, reshape, slice};
8use crate::primitive::{Primitive, ScatterReduce};
9use crate::view::normalize_axis;
10use crate::{Array, Error, Result, SliceArg, promote_types};
11
12fn single(primitive: Primitive, inputs: Vec<Array>) -> Result<Array> {
13    Ok(Array::from_op(primitive, inputs)?.swap_remove(0))
14}
15
16fn isize_shape(shape: &[usize]) -> Vec<isize> {
17    shape.iter().map(|&d| d as isize).collect()
18}
19
20/// `a` and `indices` broadcast against each other on every axis but `axis` (which keeps
21/// each one's own length), as MLX's `take_along_axis` and `put_along_axis` require.
22fn broadcast_except(
23    a: &Array,
24    indices: &Array,
25    axis: usize,
26    a_too: bool,
27) -> Result<(Array, Array)> {
28    if indices.ndim() != a.ndim() {
29        return Err(Error::InvalidArgument(format!(
30            "indices of {} dimensions for an array of {}",
31            indices.ndim(),
32            a.ndim()
33        )));
34    }
35    let mut common = Vec::with_capacity(a.ndim());
36    for (d, (&x, &i)) in a.shape().iter().zip(indices.shape()).enumerate() {
37        common.push(match (d == axis, x, i) {
38            (true, _, _) => 0,
39            (_, x, i) if x == i || i == 1 => x,
40            (_, 1, i) if a_too => i,
41            _ => {
42                return Err(Error::ShapeMismatch {
43                    op: "take_along_axis",
44                    lhs: a.shape().to_vec(),
45                    rhs: indices.shape().to_vec(),
46                });
47            }
48        });
49    }
50    let with = |len: usize| {
51        let mut s = common.clone();
52        s[axis] = len;
53        s
54    };
55    let a = if a_too {
56        broadcast_to(a, &with(a.shape()[axis]))?
57    } else {
58        a.clone()
59    };
60    let indices = broadcast_to(indices, &with(indices.shape()[axis]))?;
61    Ok((a, indices))
62}
63
64/// The elements of `a` at `indices` along `axis` (MLX's `take_along_axis`): `indices`
65/// has `a`'s number of dimensions, and both broadcast on the other axes.
66///
67/// ```
68/// use nuvec_core::{Array, Result, take_along_axis};
69///
70/// # fn main() -> Result<()> {
71/// let a = Array::from_slice(&[10.0f32, 20.0, 30.0, 40.0], &[2, 2])?;
72/// let i = Array::from_slice(&[1i32, 0], &[2, 1])?;
73/// assert_eq!(take_along_axis(&a, &i, 1)?.to_vec::<f32>()?, [20.0, 30.0]);
74/// # Ok(())
75/// # }
76/// ```
77pub fn take_along_axis(a: &Array, indices: &Array, axis: isize) -> Result<Array> {
78    let axis = normalize_axis(axis, a.ndim())?;
79    let (a, indices) = broadcast_except(a, indices, axis, true)?;
80    single(Primitive::GatherAxis { axis }, vec![a, indices])
81}
82
83/// The elements of `a` at `indices` along `axis` (MLX's `take`): the axis is replaced by
84/// `indices`' dimensions. Without an axis, `a` is flattened first.
85///
86/// ```
87/// use nuvec_core::{Array, Result, take};
88///
89/// # fn main() -> Result<()> {
90/// // An embedding lookup: rows of a table.
91/// let table = Array::from_slice(&[0.0f32, 0.5, 1.0, 1.5, 2.0, 2.5], &[3, 2])?;
92/// let tokens = Array::from_slice(&[2u32, 0, 2], &[3])?;
93/// let rows = take(&table, &tokens, Some(0))?;
94/// assert_eq!(rows.shape(), [3, 2]);
95/// assert_eq!(rows.to_vec::<f32>()?, [2.0, 2.5, 0.0, 0.5, 2.0, 2.5]);
96/// # Ok(())
97/// # }
98/// ```
99pub fn take(a: &Array, indices: &Array, axis: Option<isize>) -> Result<Array> {
100    let (a, axis) = match axis {
101        Some(axis) => (a.clone(), normalize_axis(axis, a.ndim())?),
102        None => (reshape(a, &[-1])?, 0),
103    };
104    let m = indices.size();
105    // Flat indices along `axis`, broadcast over the other axes.
106    let mut along = vec![1; a.ndim()];
107    along[axis] = m as isize;
108    let mut full = a.shape().to_vec();
109    full[axis] = m;
110    let flat = broadcast_to(&reshape(indices, &along)?, &full)?;
111    let out = single(Primitive::GatherAxis { axis }, vec![a.clone(), flat])?;
112    let mut shape = isize_shape(&a.shape()[..axis]);
113    shape.extend(isize_shape(indices.shape()));
114    shape.extend(isize_shape(&a.shape()[axis + 1..]));
115    reshape(&out, &shape)
116}
117
118fn scatter(
119    a: &Array,
120    indices: &Array,
121    values: &Array,
122    axis: isize,
123    reduce: ScatterReduce,
124) -> Result<Array> {
125    let axis = normalize_axis(axis, a.ndim())?;
126    let (_, indices) = broadcast_except(a, indices, axis, false)?;
127    let values = broadcast_to(&values.astype(a.dtype())?, indices.shape())?;
128    single(
129        Primitive::ScatterAxis { axis, reduce },
130        vec![a.clone(), indices, values],
131    )
132}
133
134/// A copy of `a` with `values` placed at `indices` along `axis` (MLX's `put_along_axis`).
135/// `indices` broadcasts against `a` on the other axes, `values` against `indices`; with
136/// repeated indices, one of the values wins.
137///
138/// ```
139/// use nuvec_core::{Array, Result, put_along_axis};
140///
141/// # fn main() -> Result<()> {
142/// let a = Array::from_slice(&[0i32; 4], &[4])?;
143/// let i = Array::from_slice(&[3i32, -4], &[2])?;
144/// let v = Array::from_slice(&[7i32, 9], &[2])?;
145/// assert_eq!(put_along_axis(&a, &i, &v, 0)?.to_vec::<i32>()?, [9, 0, 0, 7]);
146/// # Ok(())
147/// # }
148/// ```
149pub fn put_along_axis(a: &Array, indices: &Array, values: &Array, axis: isize) -> Result<Array> {
150    scatter(a, indices, values, axis, ScatterReduce::None)
151}
152
153/// A copy of `a` with `values` added at `indices` along `axis`; repeated indices
154/// accumulate (MLX's `scatter_add_axis`, the gradient of `take_along_axis`).
155pub fn scatter_add_along_axis(
156    a: &Array,
157    indices: &Array,
158    values: &Array,
159    axis: isize,
160) -> Result<Array> {
161    scatter(a, indices, values, axis, ScatterReduce::Sum)
162}
163
164/// `arrays` joined along `axis` (MLX's `concatenate`): equal shapes except along `axis`,
165/// with dtypes promoted to a common one.
166///
167/// ```
168/// use nuvec_core::{Array, Result, concatenate};
169///
170/// # fn main() -> Result<()> {
171/// let a = Array::from_slice(&[1.0f32, 2.0], &[1, 2])?;
172/// let b = Array::from_slice(&[3i32, 4, 5, 6], &[2, 2])?;
173/// let c = concatenate(&[a, b], 0)?;
174/// assert_eq!(c.shape(), [3, 2]);
175/// assert_eq!(c.to_vec::<f32>()?, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
176/// # Ok(())
177/// # }
178/// ```
179pub fn concatenate(arrays: &[Array], axis: isize) -> Result<Array> {
180    let first = arrays
181        .first()
182        .ok_or_else(|| Error::InvalidArgument("concatenate: no arrays to join".into()))?;
183    let axis = normalize_axis(axis, first.ndim())?;
184    let dtype = arrays
185        .iter()
186        .skip(1)
187        .fold(first.dtype(), |d, a| promote_types(d, a.dtype()));
188    let arrays = arrays
189        .iter()
190        .map(|a| a.astype(dtype))
191        .collect::<Result<Vec<_>>>()?;
192    single(Primitive::Concatenate { axis }, arrays)
193}
194
195/// `arrays` (one shape) stacked along a new axis `axis` (MLX's `stack`).
196pub fn stack(arrays: &[Array], axis: isize) -> Result<Array> {
197    let first = arrays
198        .first()
199        .ok_or_else(|| Error::InvalidArgument("stack: no arrays to stack".into()))?;
200    let axis = normalize_axis(axis, first.ndim() + 1)? as isize;
201    if let Some(other) = arrays.iter().find(|a| a.shape() != first.shape()) {
202        return Err(Error::ShapeMismatch {
203            op: "stack",
204            lhs: first.shape().to_vec(),
205            rhs: other.shape().to_vec(),
206        });
207    }
208    let expanded = arrays
209        .iter()
210        .map(|a| expand_dims(a, axis))
211        .collect::<Result<Vec<_>>>()?;
212    concatenate(&expanded, axis)
213}
214
215/// `a` split into `parts` equal parts along `axis` (MLX's `split` with a count); views.
216pub fn split(a: &Array, parts: usize, axis: isize) -> Result<Vec<Array>> {
217    let d = normalize_axis(axis, a.ndim())?;
218    let n = a.shape()[d];
219    if parts == 0 || !n.is_multiple_of(parts) {
220        return Err(Error::InvalidArgument(format!(
221            "split: an axis of length {n} cannot be split into {parts} equal parts"
222        )));
223    }
224    let size = n / parts;
225    let at: Vec<usize> = (1..parts).map(|k| k * size).collect();
226    split_at(a, &at, axis)
227}
228
229/// `a` split along `axis` before each of the (non-decreasing) positions `indices`, as
230/// MLX's `split` with indices: `indices.len() + 1` views.
231pub fn split_at(a: &Array, indices: &[usize], axis: isize) -> Result<Vec<Array>> {
232    let d = normalize_axis(axis, a.ndim())?;
233    let n = a.shape()[d];
234    let mut bounds = vec![0];
235    bounds.extend(indices.iter().map(|&i| i.min(n)));
236    bounds.push(n);
237    bounds
238        .windows(2)
239        .map(|w| {
240            let mut args = vec![SliceArg::full(); a.ndim()];
241            args[d] = SliceArg::new(Some(w[0] as isize), Some(w[1].max(w[0]) as isize), 1);
242            slice(a, &args)
243        })
244        .collect()
245}
nuvec_core/indexing.rs

nuvec_core/
indexing.rs
