tf.math.top_k

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

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

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

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

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

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