jax.numpy.argsort

目录

jax.numpy.argsort#

jax.numpy.argsort(a, axis=-1, *, kind=None, order=None, stable=True, descending=False)[源代码][源代码]#

返回排序数组的索引。

JAX 实现的 numpy.argsort()。

参数:

  • a (ArrayLike) – 待排序的数组

  • axis (int | None) – 沿其排序的整数轴。默认为 -1，即最后一个轴。如果为 None，则在排序前 a 会被展平。

  • stable (bool) – 指定是否应使用稳定排序的布尔值。默认=True。

  • descending (bool) – 指定是否按降序排序的布尔值。默认值=False。

  • kind (None) – 已弃用；请改用 stable=True 或 stable=False 指定排序算法。

  • order (None) – 不受 JAX 支持

返回:

对数组进行排序的索引数组。返回的数组形状将是 a.shape``（如果 ``axis 是整数）或形状 (a.size,)``（如果 ``axis 是 None）。

返回类型:

Array

示例

简单的一维排序

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

沿数组的最后一个轴排序：

>>> x = jnp.array([[2, 1, 3],
...                [6, 4, 3]])
>>> indices = jnp.argsort(x, axis=1)
>>> indices
Array([[1, 0, 2],
       [2, 1, 0]], dtype=int32)
>>> jnp.take_along_axis(x, indices, axis=1)
Array([[1, 2, 3],
       [3, 4, 6]], dtype=int32)

参见

目录