jax.numpy.flatnonzero

jax.numpy.flatnonzero(a, *, size=None, fill_value=None)

返回扁平化数组中非零元素的索引。

JAX 实现的 numpy.flatnonzero()。

jnp.flatnonzero(x) 等价于 nonzero(ravel(a))[0]。关于此函数的参数的完整讨论，请参阅 jax.numpy.nonzero()。

参数:

• a (ArrayLike) – N 维数组。

• size ( int | None | None) – 可选的静态整数，指定要返回的非零条目的数量。有关此参数的更多讨论，请参阅 jax.numpy.nonzero()。

• fill_value (None | ArrayLike | tuple[ArrayLike, ...] | None) – 指定 size 时可选的填充值。默认为 0。有关此参数的更多讨论，请参阅 jax.numpy.nonzero()。

返回:

包含扁平化数组中每个非零值的索引的数组。

返回类型:

Array

另请参阅

• jax.numpy.nonzero()

• jax.numpy.where()

示例

>>> x = jnp.array([[0, 5, 0],
...                [6, 0, 8]])
>>> jnp.flatnonzero(x)
Array([1, 3, 5], dtype=int32)

这等价于在扁平化数组上调用 nonzero()，并提取结果元组中的第一个条目。

>>> jnp.nonzero(x.ravel())[0]
Array([1, 3, 5], dtype=int32)

返回的索引可以用来从扁平化数组中提取非零条目。

>>> indices = jnp.flatnonzero(x)
>>> x.ravel()[indices]
Array([5, 6, 8], dtype=int32)