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