tf.keras.ops.argsort

Returns the indices that would sort a tensor.

tf.keras.ops.argsort(
    x, axis=-1
)

x Input tensor.
axis Axis along which to sort. Defaults to-1 (the last axis). If None, the flattened tensor is used.

Tensor of indices that sort x along the specified axis.

Examples:

One dimensional array:

x = keras.ops.array([3, 1, 2])
keras.ops.argsort(x)
array([1, 2, 0], dtype=int32)

Two-dimensional array:

>>> x = keras.ops.array([[0, 3], [3, 2], [4, 5]])
>>> x
array([[0, 3],
       [3, 2],
       [4, 5]], dtype=int32)
>>> keras.ops.argsort(x, axis=0)
array([[0, 1],
       [1, 0],
       [2, 2]], dtype=int32)
>>> keras.ops.argsort(x, axis=1)
array([[0, 1],
       [1, 0],
       [0, 1]], dtype=int32)