jax.numpy.setdiff1d

内容

jax.numpy.setdiff1d#

jax.numpy.setdiff1d(ar1, ar2, assume_unique=False, *, size=None, fill_value=None)[source]#

计算两个 1D 数组的集合差。

JAX 实现 numpy.setdiff1d().

因为 setdiff1d 输出的大小取决于数据，所以该函数通常不与 jit() 和其他 JAX 变换兼容。JAX 版本添加了可选的 size 参数，该参数必须静态指定才能在这些上下文中使用 jnp.setdiff1d。

参数：:

  • ar1 (ArrayLike) – 要进行差分的第一个元素数组。

  • ar2 (ArrayLike) – 要进行差分的第二个元素数组。

  • assume_unique (bool) – 如果为 True，则假设输入数组包含唯一的值。这允许更有效的实现，但如果 assume_unique 为 True 且输入数组包含重复项，则行为未定义。默认值：False。

  • size (int | None | None) – 如果指定，则仅返回前 size 个排序的元素。如果元素数量少于 size 指示的数量，则返回值将用 fill_value 填充。

  • fill_value (ArrayLike | None | None) – 当 size 指定且元素数量少于指示的数量时，用 fill_value 填充剩余的条目。默认为最小值。

返回:

即，ar1 中不包含在 ar2 中的元素。

返回类型:

包含输入数组中元素的集合差的数组

另请参阅

示例

计算两个数组的集合差

>>> ar1 = jnp.array([1, 2, 3, 4])
>>> ar2 = jnp.array([3, 4, 5, 6])
>>> jnp.setdiff1d(ar1, ar2)
Array([1, 2], dtype=int32)

由于输出形状是动态的，这将在 jit() 和其他转换下失败

>>> jax.jit(jnp.setdiff1d)(ar1, ar2)  
Traceback (most recent call last):
   ...
ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[4].
The error occurred while tracing the function setdiff1d at /Users/vanderplas/github/google/jax/jax/_src/numpy/setops.py:64 for jit. This concrete value was not available in Python because it depends on the value of the argument ar1.

为了确保静态已知的输出形状，您可以传递一个静态的 size 参数

>>> jit_setdiff1d = jax.jit(jnp.setdiff1d, static_argnames=['size'])
>>> jit_setdiff1d(ar1, ar2, size=2)
Array([1, 2], dtype=int32)

如果 size 太小，则差值将被截断

>>> jit_setdiff1d(ar1, ar2, size=1)
Array([1], dtype=int32)

如果 size 太大，则输出将用 fill_value 填充

>>> jit_setdiff1d(ar1, ar2, size=4, fill_value=0)
Array([1, 2, 0, 0], dtype=int32)
内容