take_along_axis

take_along_axis(x: array, indices: array, /, *, axis: int = -1) → array

Returns elements from an array at the one-dimensional indices specified by indices along a provided axis.

Parameters:

  • x (array) – input array. Must be compatible with indices, except for the axis specified by axis (see Broadcasting).

  • indices (array) – array indices. Must have the same number of axes as x and must be compatible with x, except for the axis specified by axis (see Broadcasting). If an index is negative, the function must determine the element to select along a specified axis by counting from the last element (where -1 refers to the last element).

  • axis (int) – axis along which to select values. A valid axis must be an integer on the interval [-N, N), where N is the number of axes in x. If an axis is specified as a negative integer, the function must determine the axis along which to perform the operation by counting backward from the last axis (where -1 refers to the last axis). If provided an invalid axis, the function must raise an exception. Default: -1.

Returns:

out (array) – an array containing elements from x. The returned array must have the same data type as x. The returned array must have the same number of axes as x and must have a shape determined according to Broadcasting, except for the axis specified by axis whose size must equal the size of the corresponding axis in indices.

Notes

  • This specification does not require bounds checking. The behavior for out-of-bounds indices is unspecified and thus implementation-defined.

New in version 2024.12.