tf.math.top_k

Finds values and indices of the k largest entries for the last dimension.

Used in the notebooks

Used in the tutorials

If the input is a vector (rank=1), finds the k largest entries in the vector and outputs their values and indices as vectors. Thus values[j] is the j-th largest entry in input, and its index is indices[j].

result = tf.math.top_k([1, 2, 98, 1, 1, 99, 3, 1, 3, 96, 4, 1],
                        k=3)
result.values.numpy()
array([99, 98, 96], dtype=int32)
result.indices.numpy()
array([5, 2, 9], dtype=int32)

For matrices (resp. higher rank input), computes the top k entries in each row (resp. vector along the last dimension). Thus,