pub fn take(a: &Array, indices: &Array, axis: Option<isize>) -> Result<Array>Expand description
The elements of a at indices along axis (MLX’s take): the axis is replaced by
indices’ dimensions. Without an axis, a is flattened first.
use nuvec_core::{Array, Result, take};
// An embedding lookup: rows of a table.
let table = Array::from_slice(&[0.0f32, 0.5, 1.0, 1.5, 2.0, 2.5], &[3, 2])?;
let tokens = Array::from_slice(&[2u32, 0, 2], &[3])?;
let rows = take(&table, &tokens, Some(0))?;
assert_eq!(rows.shape(), [3, 2]);
assert_eq!(rows.to_vec::<f32>()?, [2.0, 2.5, 0.0, 0.5, 2.0, 2.5]);