Save the date! Google I/O returns May 18-20 Register now

tfp.substrates.jax.stats.lexicographical_indirect_sort

Sorts by primary, then by secondary returning the indices.

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.