mlx.core.take

take(a: array, /, indices: int | array, axis: int | None = None, *, stream: None | Stream | Device = None) → array

Take elements along an axis.

The elements are taken from indices along the specified axis. If the axis is not specified the array is treated as a flattened 1-D array prior to performing the take.

As an example, if the axis=1 this is equivalent to a[:, indices, ...].

Parameters:

• a (array) – Input array.

• indices (int or array) – Integer index or input array with integral type.

• axis (int, optional) – Axis along which to perform the take. If unspecified the array is treated as a flattened 1-D vector.

Returns:

The indexed values of a.

Return type:

array