jax.numpy.lexsort

jax.numpy.lexsort(keys, axis=-1)

Perform an indirect stable sort using a sequence of keys.

LAX-backend implementation of numpy.lexsort().

Original docstring below.

Given multiple sorting keys, which can be interpreted as columns in a spreadsheet, lexsort returns an array of integer indices that describes the sort order by multiple columns. The last key in the sequence is used for the primary sort order, the second-to-last key for the secondary sort order, and so on. The keys argument must be a sequence of objects that can be converted to arrays of the same shape. If a 2D array is provided for the keys argument, its rows are interpreted as the sorting keys and sorting is according to the last row, second last row etc.

Parameters:

• keys ((k, N) array or tuple containing k (N,)-shaped sequences) – The k different “columns” to be sorted. The last column (or row if keys is a 2D array) is the primary sort key.

• axis (int, optional) – Axis to be indirectly sorted. By default, sort over the last axis.

Returns:

indices – Array of indices that sort the keys along the specified axis.

Return type:

(N,) ndarray of ints