jax.numpy.flatnonzero

目录

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) – 可选的静态整数，指定要返回的非零条目的数量。有关此参数的更多讨论，请参见 jax.numpy.nonzero()。

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

返回:

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

返回类型:

Array

参见

示例

>>> 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)
目录