jax.numpy.bincount#

jax.numpy.bincount(x, weights=None, minlength=0, *, length=None)[来源]#

计算整数数组中每个值出现的次数。

JAX 实现 numpy.bincount()。

对于一个非负整数数组 x，此函数返回一个大小为 x.max() + 1 的数组 counts，其中 counts[i] 包含值 i 在 x 中的出现次数。

JAX 版本与 NumPy 版本有几处不同。

  • 在 NumPy 中，传递一个带有负数的数组 x 会导致错误。在 JAX 中，负值会被裁剪为零。

  • JAX 添加了一个可选的 length 参数，可用于静态指定输出数组的长度，以便此函数可与 jax.jit() 等转换一起使用。在这种情况下，大于 length + 1 的项将被丢弃。

参数:

  • x (ArrayLike) – 非负整数的一维数组。

  • weights (ArrayLike | None) – 与 x 相关的可选权重数组。如果未指定，则每个条目的权重将是 1。

  • minlength (int) – 输出 counts 数组的最小长度。

  • length (int | None) – 输出 counts 数组的长度。必须静态指定 length 才能将 bincount 与 jax.jit() 和其他 JAX 转换一起使用。

返回:

一个 counts 数组或加权后的 counts 数组，反映了 x 中值的出现次数。

返回类型:

Array

另请参阅

示例

基本 bincount

>>> x = jnp.array([1, 1, 2, 3, 3, 3])
>>> jnp.bincount(x)
Array([0, 2, 1, 3], dtype=int32)

加权 bincount

>>> weights = jnp.array([1, 2, 3, 4, 5, 6])
>>> jnp.bincount(x, weights)
Array([ 0,  3,  3, 15], dtype=int32)

指定静态 length 使其与 jit 兼容。

>>> jit_bincount = jax.jit(jnp.bincount, static_argnames=['length'])
>>> jit_bincount(x, length=5)
Array([0, 2, 1, 3, 0], dtype=int32)

任何负数都会被裁剪到第一个 bin，超出指定 length 的数字将被丢弃。

>>> x = jnp.array([-1, -1, 1, 3, 10])
>>> jnp.bincount(x, length=5)
Array([2, 1, 0, 1, 0], dtype=int32)