tf.keras.ops.top_k

See Stable See Nightly

Finds the top-k values and their indices in a tensor.

tf.keras.ops.top_k(
    x, k, sorted=True
)

Args
x	Input tensor.
k	An integer representing the number of top elements to retrieve.
sorted	A boolean indicating whether to sort the output in descending order. Defaults to True.

Returns
A tuple containing two tensors. The first tensor contains the top-k values, and the second tensor contains the indices of the top-k values in the input tensor.

Example:

x = keras.ops.convert_to_tensor([5, 2, 7, 1, 9, 3])
values, indices = top_k(x, k=3)
print(values)
array([9 7 5], shape=(3,), dtype=int32)
print(indices)
array([4 2 0], shape=(3,), dtype=int32)