jax.lax.top_k

Contents

jax.lax.top_k#

jax.lax.top_k(operand, k, *, axis=-1)[source]#

Returns top k values and their indices along the specified axis of operand.

Parameters:

  • operand (ArrayLike) – N-dimensional array of non-complex type.

  • k (int) – integer specifying the number of top entries.

  • axis (int) – optional integer specifying the axis along which to compute the top k entries. Default is -1, indicating the last axis.

Returns:

A tuple (values, indices) where

  • values is an array containing the top k values along the last axis.

  • indices is an array containing the indices corresponding to values.

Return type:

tuple[Array, Array]

values[..., i, ...] is the i-th largest entry in operand along the specified axis, and its index is indices[..., i, ...].

If two elements are equal, the lower-index element appears first.

See also

Examples

Find the largest three values, and their indices, within an array:

>>> x = jnp.array([9., 3., 6., 4., 10.])
>>> values, indices = jax.lax.top_k(x, 3)
>>> values
Array([10.,  9.,  6.], dtype=float32)
>>> indices
Array([4, 0, 2], dtype=int32)
Contents