pub fn take_along_axis(a: &Array, indices: &Array, axis: isize) -> Result<Array>Expand description
The elements of a at indices along axis (MLX’s take_along_axis): indices
has a’s number of dimensions, and both broadcast on the other axes.
use nuvec_core::{Array, Result, take_along_axis};
let a = Array::from_slice(&[10.0f32, 20.0, 30.0, 40.0], &[2, 2])?;
let i = Array::from_slice(&[1i32, 0], &[2, 1])?;
assert_eq!(take_along_axis(&a, &i, 1)?.to_vec::<f32>()?, [20.0, 30.0]);