mlx.core.split#

split(a: array, /, indices_or_sections: int | Sequence[int], axis: int = 0, *, stream: None | Stream | Device = None) → array#

Split an array along a given axis.

Parameters:

  • a (array) – Input array.

  • indices_or_sections (int or list(int)) – If indices_or_sections is an integer the array is split into that many sections of equal size. An error is raised if this is not possible. If indices_or_sections is a list, the list contains the indices of the start of each subarray along the given axis.

  • axis (int, optional) – Axis to split along, defaults to 0.

Returns:

A list of split arrays.

Return type:

list(array)