Skip to main content

put_along_axis

Function put_along_axis 

Source
pub fn put_along_axis(
    a: &Array,
    indices: &Array,
    values: &Array,
    axis: isize,
) -> Result<Array>
Expand description

A copy of a with values placed at indices along axis (MLX’s put_along_axis). indices broadcasts against a on the other axes, values against indices; with repeated indices, one of the values wins.

use nuvec_core::{Array, Result, put_along_axis};

let a = Array::from_slice(&[0i32; 4], &[4])?;
let i = Array::from_slice(&[3i32, -4], &[2])?;
let v = Array::from_slice(&[7i32, 9], &[2])?;
assert_eq!(put_along_axis(&a, &i, &v, 0)?.to_vec::<i32>()?, [9, 0, 0, 7]);