jax.numpy.lexsort

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

按字典顺序对键序列进行排序。

JAX 对 numpy.lexsort() 的实现。

参数:

• keys (Array | ndarray | Sequence[Array | ndarray | bool | number | bool | int | float | complex]) – 要排序的数组序列；所有数组必须具有相同的形状。序列中的最后一个键用作主键。

• axis (int) – 沿其排序的轴（默认值：-1）。

返回:

一个整数数组，形状为 keys[0].shape，给出按字典顺序排序的条目的索引。

返回类型:

Array

另请参阅

• jax.numpy.argsort()：按索引对单个条目进行排序。

• jax.lax.sort()：直接 XLA 排序 API。

示例

具有单个键的 lexsort() 等效于 argsort()

>>> key1 = jnp.array([4, 2, 3, 2, 5])
>>> jnp.lexsort([key1])
Array([1, 3, 2, 0, 4], dtype=int32)
>>> jnp.argsort(key1)
Array([1, 3, 2, 0, 4], dtype=int32)

对于多个键，lexsort() 使用最后一个键作为主键

>>> key2 = jnp.array([2, 1, 1, 2, 2])
>>> jnp.lexsort([key1, key2])
Array([1, 2, 3, 0, 4], dtype=int32)

当打印排序后的键时，索引的含义变得更加清晰

>>> indices = jnp.lexsort([key1, key2])
>>> print(f"{key1[indices]}\n{key2[indices]}")
[2 3 2 4 5]
[1 1 2 2 2]

请注意，key2 的元素按顺序出现，并且在重复值序列中，`key1 的对应元素也按顺序出现。

对于多维输入，lexsort() 默认沿最后一个轴排序

>>> key1 = jnp.array([[2, 4, 2, 3],
...                   [3, 1, 2, 2]])
>>> key2 = jnp.array([[1, 2, 1, 3],
...                   [2, 1, 2, 1]])
>>> jnp.lexsort([key1, key2])
Array([[0, 2, 1, 3],
       [1, 3, 2, 0]], dtype=int32)

可以使用 axis 关键字选择不同的排序轴；这里我们沿前导轴排序

>>> jnp.lexsort([key1, key2], axis=0)
Array([[0, 1, 0, 1],
       [1, 0, 1, 0]], dtype=int32)