top_k

Contents

top_k#

class saiunit.lax.top_k(operand, k, **kwargs)#

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

Parameters:

  • operand (saiunit.Quantity | Array | ndarray | bool | number | bool | int | float | complex) – N-dimensional array of non-complex type.

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

Return type:

tuple[saiunit.Quantity | Array, Array | ndarray | bool | number | bool | int | float | complex]

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.

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