NonZero

Returns a tensor containing the indices of all non-zero elements of input.

Abstract Signature:

NonZero(input: Tensor)

PyTorch

API: torch.nonzero
Strategy: Macro 'torch.nonzero({a}, as_tuple=True)'

JAX (Core)

API: jax.numpy.argwhere
Strategy: Direct Mapping
Official Docs ↗

NumPy

API: numpy.argwhere
Strategy: Direct Mapping
Official Docs ↗

Keras

API: keras.ops.nonzero
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

API: tf.where
Strategy: Direct Mapping
Official Docs ↗

Apple MLX

API: mlx.core.argwhere
Strategy: Direct Mapping
Official Docs ↗

Flax NNX

API: jax.numpy.argwhere
Strategy: Direct Mapping
Official Docs ↗

PaxML / Praxis

API: jax.numpy.argwhere
Strategy: Direct Mapping
Official Docs ↗