tfp.stats.lexicographical_indirect_sort

Sorts by primary, then by secondary returning the indices.

tfp.stats.lexicographical_indirect_sort(
    primary, secondary, name=None
)

primary a Tensor of shape [n] containing the primary sort key. the primary sort key value.
secondary a Tensor of shape [n] containing the secondary sort key to be used when the primary keys are identical.
name Optional Python str name for ops created by this method. Default value: None (i.e., 'lexicographical_indirect_sort').

lexicographic A permutation of range(n) that provides the sorted primary, then secondary values.