flatnonzero

Contents

flatnonzero#

class brainunit.math.flatnonzero(a, *, size=None, fill_value=None, **kwargs)#

Return indices that are non-zero in the flattened input.

Units are stripped before the search.

Parameters:

  • a (Array | ndarray | bool | number | bool | int | float | complex | saiunit.Quantity) – Input data.

  • size (int | None) – Fixed output size (for use inside jax.jit).

  • fill_value (Array | ndarray | bool | number | bool | int | float | complex | None) – Fill value for padding when size is given.

Returns:

indices – Indices of non-zero elements in the flattened array.

Return type:

Array

Examples

>>> import saiunit as u
>>> import jax.numpy as jnp
>>> u.math.flatnonzero(jnp.array([0, 1, 0, 2]), size=2)
Array([1, 3], dtype=int32)
Contents