Function split_equal

pub fn split_equal(
    a: impl AsRef<Array>,
    num_parts: i32,
    axis: impl Into<Option<i32>>,
) -> Result<Vec<Array>>

Split an array into equal parts along a given axis. Returns an error if the array cannot be split into equal parts.

Params

• a: The input array.
• num_parts: The number of parts to split into.
• axis: The axis to split along. Default is 0 if not provided.

Example

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

let a = Array::from_iter(0..10, &[10]);
let result = split_equal(&a, 2, 0);