jax.numpy.flatnonzero

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

返回展平数组中非零元素的索引

numpy.flatnonzero() 的 JAX 实现。

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()。

返回:

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

返回类型:

数组

另请参阅

• 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)