FlatnonzeroΒΆ

Return indices that are non-zero in the flattened version of a.

Abstract Signature:

Flatnonzero(a: Array)

PyTorch

API: β€”
Strategy: Macro 'torch.nonzero({a}.flatten(), as_tuple=True)[0]'

JAX (Core)

API: jax.numpy.flatnonzero
Strategy: Direct Mapping

NumPy

API: numpy.flatnonzero
Strategy: Direct Mapping

Flax NNX

API: jax.numpy.flatnonzero
Strategy: Direct Mapping

PaxML / Praxis

API: jax.numpy.flatnonzero
Strategy: Direct Mapping