jax.numpy.bincount

内容

jax.numpy.bincount#

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

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

numpy.bincount() 的 JAX 实现。

对于正整数数组 x，此函数返回大小为 x.max() + 1 的数组 counts，使得 counts[i] 包含 x 中值 i 的出现次数。

JAX 版本与 NumPy 版本有一些区别

  • 在 NumPy 中，传递包含负数的数组 x 将导致错误。在 JAX 中，负值会被截断为零。

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

参数：

  • x (ArrayLike) – 正整数的 N 维数组

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

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

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

返回值：

一个反映 x 中值的出现次数的计数或加权总和数组。

返回类型：

数组

另请参阅

示例

基本 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)
目录