dpnp.take_along_axis

dpnp.take_along_axis(a, indices, axis)

Take values from the input array by matching 1d index and data slices.

For full documentation refer to numpy.take_along_axis.

Parameters:

• a ({dpnp.ndarray, usm_ndarray}, (Ni..., M, Nk...)) – Source array

• indices ({dpnp.ndarray, usm_ndarray}, (Ni..., J, Nk...)) – Indices to take along each 1d slice of a. This must match the dimension of the input array, but dimensions Ni and Nj only need to broadcast against a.

• axis (int) – The axis to take 1d slices along. If axis is None, the input array is treated as if it had first been flattened to 1d, for consistency with sort and argsort.

Returns:

out – The indexed result.

Return type:

dpnp.ndarray

See also

dpnp.take

Take along an axis, using the same indices for every 1d slice.

dpnp.put_along_axis

Put values into the destination array by matching 1d index and data slices.

dpnp.argsort

Return the indices that would sort an array.

Examples

For this sample array

>>> import dpnp as np
>>> a = np.array([[10, 30, 20], [60, 40, 50]])

We can sort either by using sort directly, or argsort and this function

>>> np.sort(a, axis=1)
array([[10, 20, 30],
       [40, 50, 60]])
>>> ai = np.argsort(a, axis=1)
>>> ai
array([[0, 2, 1],
       [1, 2, 0]])
>>> np.take_along_axis(a, ai, axis=1)
array([[10, 20, 30],
       [40, 50, 60]])

The same works for max and min, if you maintain the trivial dimension with keepdims:

>>> np.max(a, axis=1, keepdims=True)
array([[30],
       [60]])
>>> ai = np.argmax(a, axis=1, keepdims=True)
>>> ai
array([[1],
       [0]])
>>> np.take_along_axis(a, ai, axis=1)
array([[30],
       [60]])

If we want to get the max and min at the same time, we can stack the indices first

>>> ai_min = np.argmin(a, axis=1, keepdims=True)
>>> ai_max = np.argmax(a, axis=1, keepdims=True)
>>> ai = np.concatenate([ai_min, ai_max], axis=1)
>>> ai
array([[0, 1],
       [1, 0]])
>>> np.take_along_axis(a, ai, axis=1)
array([[10, 30],
       [40, 60]])