tf.math.top_k

Находит значения и индексы k наибольших элементов по последнему измерению.

tf.math.top_k(
    input, k=1, sorted=True, name=None
)

Если входной параметр является вектором (ранг=1), находит k наибольших элементов в векторе и выводит их значения и индексы в виде векторов. Таким образом, values[j] — это j-й по величине элемент в input, а его индекс равен 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)

Для матриц (соответственно, входных данных более высокого ранга) вычисляет k наибольших элементов в каждом ряду (соответственно, векторе по последнему измерению). Таким образом,

input = tf.random.normal(shape=(3,4,5,6))
k = 2
values, indices  = tf.math.top_k(input, k=k)
values.shape.as_list()
[3, 4, 5, 2]

values.shape == indices.shape == input.shape[:-1] + [k]
True

Индексы могут быть использованы для gather из тензора, чья форма соответствует input.

gathered_values = tf.gather(input, indices, batch_dims=-1)
assert tf.reduce_all(gathered_values == values)

Если два элемента равны, элемент с меньшим индексом появляется первым.

result = tf.math.top_k([1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0],
                       k=3)
result.indices.numpy()
array([0, 1, 3], dtype=int32)