dpctl.tensor.top_k

dpctl.tensor.top_k(x, k, axis=None, mode='largest')[source]

Returns the k largest or smallest values and their indices in the input array x along the specified axis axis.

Parameters:
  • x (usm_ndarray) – input array.

  • k (int) – number of elements to find. Must be a positive integer value.

  • axis (Optional[int]) – axis along which to search. If None, the search will be performed over the flattened array. Default: None.

  • mode (Literal["largest", "smallest"]) –

    search mode. Must be one of the following modes:

    • ”largest”: return the k largest elements.

    • ”smallest”: return the k smallest elements.

    Default: “largest”.

Returns:

tuple[usm_ndarray, usm_ndarray]

a namedtuple (values, indices) whose

  • first element values will be an array containing the k largest or smallest elements of x. The array has the same data type as x. If axis was None, values will be a one-dimensional array with shape (k,) and otherwise, values will have shape x.shape[:axis] + (k,) + x.shape[axis+1:]

  • second element indices will be an array containing indices of x that result in values. The array will have the same shape as values and will have the default array index data type.