jax.numpy.top_k

Contents

jax.numpy.top_k#

jax.numpy.top_k(a, k, /, *, axis=-1, mode='largest', sorted=True)[source]#

Return the k largest or smallest elements and their indices along an axis.

JAX implementation of numpy.top_k().

Parameters:

  • a (Array | ndarray | bool | number | bool | int | float | complex) – array to find top k elements from.

  • k (int) – static integer specifying the number of elements to return. Must be a non-negative integer and no larger than the size of the array along the specified axis.

  • axis (int) – static integer axis along which to find top k elements. Default is -1.

  • mode (str) – string specifying whether to return the 'largest' (default) or 'smallest' elements.

  • sorted (bool) – boolean specifying whether to return elements in sorted order. Default is True.

Returns:

A tuple of (topk_values, topk_indices), where topk_values are the top k values and topk_indices are the corresponding indices. Both arrays have the shape of a with the dimension along axis replaced by k.

Return type:

tuple[Array, Array]

See also

Examples

Find the two largest elements along the last axis:

>>> a = jnp.array([[1, 2, 3, 4, 5],
...                [5, 4, 3, 2, 1]])
>>> values, indices = jnp.top_k(a, 2)
>>> values
Array([[5, 4],
       [5, 4]], dtype=int32)
>>> indices
Array([[4, 3],
       [0, 1]], dtype=int32)

Find the two smallest elements along the first axis:

>>> values, indices = jnp.top_k(a, 2, axis=0, mode='smallest')
>>> values
Array([[1, 2, 3, 2, 1],
       [5, 4, 3, 4, 5]], dtype=int32)
>>> indices
Array([[0, 0, 0, 1, 1],
       [1, 1, 1, 0, 0]], dtype=int32)
Contents