jax.numpy.argsort#
- jax.numpy.argsort(a, axis=-1, *, kind=None, order=None, stable=True, descending=False)[源代码]#
返回对数组进行排序的索引。
numpy.argsort()
的 JAX 实现。- 参数:
- 返回:
返回一个数组的索引,该索引对数组进行排序。返回的数组的形状将为
a.shape
(如果axis
是一个整数)或(a.size,)
(如果axis
为 None)。- 返回类型:
示例
简单的一维排序
>>> 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)
另请参阅
jax.numpy.sort()
:直接返回排序后的值。jax.numpy.lexsort()
:对多个数组进行字典式排序。jax.lax.sort()
:包装 XLA Sort 操作符的底层函数。