dpnp.tensor.take_along_axis

dpnp.tensor.take_along_axis(x, indices, /, *, axis=-1, mode='wrap')

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

Parameters:

• x (usm_ndarray) -- input array. Must be compatible with indices, except for the axis (dimension) specified by axis.

• indices (usm_ndarray) -- array indices. Must have the same rank (i.e., number of dimensions) as x.

• axis (int, optional) --

  axis along which to select values. If axis is negative, the function determines the axis along which to select values by counting from the last dimension.

  Default: -1.

• mode ({"wrap", "clip"}, optional) --

  How out-of-bounds indices will be handled. Possible values are:

  • "wrap": clamps indices to (-n <= i < n), then wraps negative indices.

  • "clip": clips indices to (0 <= i < n).

  Default: "wrap".

Returns:

out -- an array having the same data type as x. The returned array has the same rank (i.e., number of dimensions) as x and a shape determined according to broadcasting rules, except for the axis (dimension) specified by axis whose size must equal the size of the corresponding axis (dimension) in indices.

Return type:

usm_ndarray

Notes

Treatment of the out-of-bound indices in indices array is controlled by the value of mode keyword.