LexSort

Perform an indirect stable sort using a sequence of keys.

Abstract Signature:

LexSort(keys: Tensor, axis: int = -1)

JAX (Core)

API: jax.numpy.lexsort
Strategy: Direct Mapping
Official Docs ↗

NumPy

API: numpy.lexsort
Strategy: Direct Mapping
Official Docs ↗

Flax NNX

API: jax.numpy.lexsort
Strategy: Direct Mapping
Official Docs ↗

PaxML / Praxis

API: jax.numpy.lexsort
Strategy: Direct Mapping
Official Docs ↗