nonzero

Contents

nonzero#

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

Return the indices of non-zero elements.

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 – Tuple of arrays, one per dimension, containing the indices of non-zero elements.

Return type:

Sequence[Array]

Examples

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