jax.numpy.where

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

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

JAX 中对 numpy.where() 的实现。

注意

当只提供 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()。

返回:

一个 dtype 为 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()（反向模式微分）计算梯度时，x 或 y 中的 NaN 将传播到梯度中，无论 condition 的值如何。有关此行为和变通方法的更多信息，请参见 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)