Function reshape_device

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

Reshape an array while preserving the size. Returns an error if the new shape is invalid.

Params

• a: The input array.
• shape: New shape.

Example

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

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