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)