Function swap_axes_device

pub fn swap_axes_device(
    a: impl AsRef<Array>,
    axis1: i32,
    axis2: i32,
    stream: impl AsRef<Stream>,
) -> Result<Array>

Swap two axes of an array. Returns an error if the axes are invalid.

Params

• a: The input array.
• axis1: The first axis.
• axis2: The second axis.

Example

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

let a = Array::from_iter(0..6, &[2, 3]);
let result = swap_axes(&a, 0, 1);