jax.numpy.argpartition

jax.numpy.argpartition#

jax.numpy.argpartition(a, kth, axis=-1)[源代码][源代码]#

返回部分排序数组的索引。

JAX 实现的 numpy.argpartition()。JAX 版本与 NumPy 在处理 NaN 条目时有所不同：设置了负位的 NaN 会被排序到数组的开头。

参数:

  • a (ArrayLike) – 要分区的数组。

  • kth (int) – 静态整数索引，用于对数组进行分区。

  • axis (int) – 静态整数轴，用于对数组进行分区；默认值为 -1。

返回:

沿 axis 在 kth 值处分割 a 的索引。kth 之前的条目是值小于 take(a, kth, axis) 的索引，kth 之后的条目是值大于 take(a, kth, axis) 的索引。

返回类型:

Array

备注

JAX 版本要求 kth 参数为静态整数，而不是通用数组。这是通过两次调用 jax.lax.top_k() 实现的。如果你只访问输出的前 k 个或后 k 个值，直接调用 jax.lax.top_k() 可能会更高效。

参见

示例

>>> x = jnp.array([6, 8, 4, 3, 1, 9, 7, 5, 2, 3])
>>> kth = 4
>>> idx = jnp.argpartition(x, kth)
>>> idx
Array([4, 8, 3, 9, 2, 0, 1, 5, 6, 7], dtype=int32)

结果是一个部分排序的索引序列。所有在 kth 之前的索引对应的值都小于枢轴值，而所有在 kth 之后的索引对应的值都大于枢轴值：

>>> x_partitioned = x[idx]
>>> smallest_values = x_partitioned[:kth]
>>> pivot_value = x_partitioned[kth]
>>> largest_values = x_partitioned[kth + 1:]
>>> print(smallest_values, pivot_value, largest_values)
[1 2 3 3] 4 [6 8 9 7 5]

注意，在 smallest_values 和 largest_values 中，返回的顺序是任意的，并且依赖于实现。

目录