Function expand_dims_device

pub fn expand_dims_device(
    a: impl AsRef<Array>,
    axes: &[i32],
    stream: impl AsRef<Stream>,
) -> Result<Array>

Add a size one dimension at the given axis, returns an error if the axes are invalid.

Params

• a: The input array.
• axes: The index of the inserted dimensions.

Example

use mlx_rs::{Array, ops::*};

let x = Array::zeros::<i32>(&[2, 2]).unwrap();
let result = expand_dims(&x, &[0]);