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 (dimension) specified by axis (see Broadcasting).

  • indices (array) – array indices. Must have the same rank (i.e., number of dimensions) as x. If an index is negative, the function must determine the element to select along a specified axis (dimension) by counting from the last element (where -1 refers to the last element).

  • axis (int) – axis along which to select values. If axis is negative, the function must determine the axis along which to select values by counting from the last dimension (where -1 refers to the last dimension). Default: -1.

Returns:

out (array) – an array having the same data type as x. Must have the same rank (i.e., number of dimensions) as x and must have a shape determined according to Broadcasting, except for the axis (dimension) specified by axis whose size must equal the size of the corresponding axis (dimension) in indices.

Notes

  • This specification does not require bounds checking. The behavior for out-of-bounds indices is left unspecified.