flatnonzero

Contents

flatnonzero#

class saiunit.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 (Union[Array, ndarray, number, bool, saiunit.Quantity]) – Input data.

  • size (Optional[int]) – Fixed output size (for use inside jax.jit).

  • fill_value (Union[Array, ndarray, number, bool, 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