dpctl.tensor.take_along_axis

dpctl.tensor.take_along_axis(x, indices, /, *, axis=-1, mode='wrap')[source]

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 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 (str, 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:

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

Note

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