tf.argsort

Возвращает индексы тензора, которые дают его отсортированный порядок вдоль оси.

tf.argsort(
    values, axis=-1, direction='ASCENDING', stable=False, name=None
)

values = [1, 10, 26.9, 2.8, 166.32, 62.3]
sort_order = tf.argsort(values)
sort_order.numpy()
array([0, 3, 1, 2, 5, 4], dtype=int32)

Для одномерного тензора:

sorted = tf.gather(values, sort_order)
assert tf.reduce_all(sorted == tf.sort(values))

Для тензоров более высокой размерности вывод имеет такую же форму, как values, но вдоль заданной оси значения представляют индекс отсортированного элемента в этом срезе тензора в заданной позиции.

mat = [[30,20,10],
       [20,10,30],
       [10,30,20]]
indices = tf.argsort(mat)
indices.numpy()
array([[2, 1, 0],
       [1, 0, 2],
       [0, 2, 1]], dtype=int32)

Если axis=-1 эти индексы могут быть использованы для применения сортировки с помощью tf.gather:

tf.gather(mat, indices, batch_dims=-1).numpy()
array([[10, 20, 30],
       [10, 20, 30],
       [10, 20, 30]], dtype=int32)

См. также:

• tf.sort: Сортировка вдоль оси.
• tf.math.top_k: Частичная сортировка, возвращающая фиксированное количество верхних значений и соответствующих индексов.