jax.numpy.where

jax.numpy.where(condition, x=None, y=None, /, *, size=None, fill_value=None)

根据条件从两个数组中选择元素。

numpy.where() 的 JAX 实现。

注意

当仅提供 condition 时，jnp.where(condition) 等效于 jnp.nonzero(condition)。对于这种情况，请参阅 jax.numpy.nonzero() 的文档。以下文档字符串侧重于指定了 x 和 y 的情况。

三项版本的 jnp.where 降低为 jax.lax.select()。

参数:

• condition – 布尔数组。当指定 x 和 y 时，必须与它们广播兼容。

• x – 类数组。应与 condition 和 y 广播兼容，并与 y 类型转换兼容。

• y – 类数组。应与 condition 和 x 广播兼容，并与 x 类型转换兼容。

• size – 整数，仅当 x 和 y 为 None 时引用。有关详细信息，请参阅 jax.numpy.nonzero()。

• fill_value – 仅当 x 和 y 为 None 时引用。有关详细信息，请参阅 jax.numpy.nonzero()。

返回:

一个数据类型为 jnp.result_type(x, y) 的数组，其值从 condition 为 True 时的 x 中提取，从 condition 为 False 时的 y 中提取。如果 x 和 y 为 None，则此函数的行为有所不同；有关返回类型的描述，请参阅 jax.numpy.nonzero()。

另请参阅

• jax.numpy.nonzero()

• jax.numpy.argwhere()

• jax.lax.select()

笔记

当 jax.numpy.where() 的 x 或 y 输入可能具有 NaN 值时，需要特别注意。具体来说，当使用 jax.grad()（反向模式微分）求梯度时，无论 condition 的值如何，x 或 y 中的 NaN 都将传播到梯度中。有关此行为和解决方法的更多信息，请参见 JAX FAQ。

示例

当未提供 x 和 y 时，where 的行为等效于 jax.numpy.nonzero()

>>> x = jnp.arange(10)
>>> jnp.where(x > 4)
(Array([5, 6, 7, 8, 9], dtype=int32),)
>>> jnp.nonzero(x > 4)
(Array([5, 6, 7, 8, 9], dtype=int32),)

当提供 x 和 y 时，where 根据指定的条件在它们之间进行选择

>>> jnp.where(x > 4, x, 0)
Array([0, 0, 0, 0, 0, 5, 6, 7, 8, 9], dtype=int32)