feat: add nonzero_numpy operator with tests and benchmark (#1755)

Co-authored-by: factnn <1050552884@qq.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: 3 people

benchmark/test_nonzero_numpy.py

@@ -0,0 +1,15 @@
+import pytest
+import torch
+
+from . import base, consts
+
+
+@pytest.mark.nonzero_numpy
+def test_nonzero_numpy():
+    bench = base.GenericBenchmark2DOnly(
+        input_fn=base.unary_input_fn,
+        op_name="nonzero_numpy",
+        torch_op=torch.ops.aten.nonzero_numpy,
+        dtypes=consts.FLOAT_DTYPES + consts.INT_DTYPES + consts.BOOL_DTYPES,
+    )
+    bench.run()

src/flag_gems/__init__.py

@@ -358,6 +358,7 @@ def torch_ge(v):
     ("nll_loss2d_backward", nll_loss2d_backward),
     ("nll_loss2d_forward", nll_loss2d_forward),
     ("nonzero", nonzero),
+    ("nonzero_numpy", nonzero_numpy),
     ("normal.Tensor_float", normal_tensor_float),
     ("normal.Tensor_Tensor", normal_tensor_tensor),
     ("normal.float_Tensor", normal_float_tensor),

src/flag_gems/ops/__init__.py

@@ -230,6 +230,7 @@
     nll_loss_forward,
 )
 from flag_gems.ops.nonzero import nonzero
+from flag_gems.ops.nonzero_numpy import nonzero_numpy
 from flag_gems.ops.normal import (
     normal_,
     normal_float_tensor,
@@ -651,6 +652,7 @@
     "nll_loss_nd_forward",
     "nll_loss_nd_backward",
     "nonzero",
+    "nonzero_numpy",
     "normal_float_tensor",
     "normal_tensor_float",
     "normal_tensor_tensor",

src/flag_gems/ops/nonzero_numpy.py

@@ -0,0 +1,22 @@
+import logging
+
+from flag_gems.ops.nonzero import nonzero
+
+logger = logging.getLogger(__name__)
+
+
+def nonzero_numpy(inp):
+    """
+    Returns a tuple of 1D tensors, one for each dimension of the input,
+    containing the indices of the non-zero elements in that dimension.
+
+    This is equivalent to torch.nonzero(...).T or numpy.nonzero() behavior.
+    """
+    logger.debug("GEMS NONZERO_NUMPY")
+
+    # Use the existing nonzero implementation which returns shape [N, ndim]
+    out = nonzero(inp, as_tuple=False)
+
+    # Unbind along dim=1 to get ndim tensors of shape [N]
+    # Convert to list since aten::nonzero_numpy returns Tensor[]
+    return list(out.unbind(dim=1))

tests/test_nonzero_numpy.py

@@ -0,0 +1,38 @@
+import pytest
+import torch
+
+import flag_gems
+
+from .accuracy_utils import (
+    BOOL_TYPES,
+    FLOAT_DTYPES,
+    INT_DTYPES,
+    REDUCTION_SHAPES,
+    gems_assert_equal,
+    to_reference,
+)
+
+NONZERO_SHAPES = REDUCTION_SHAPES
+
+
+@pytest.mark.nonzero_numpy
+@pytest.mark.parametrize("shape", NONZERO_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES + INT_DTYPES + BOOL_TYPES)
+def test_nonzero_numpy(shape, dtype):
+    if dtype == torch.bool:
+        inp = torch.randint(0, 2, shape, dtype=torch.int, device=flag_gems.device).to(
+            dtype
+        )
+    elif dtype in INT_DTYPES:
+        inp = torch.randint(-3, 3, shape, device=flag_gems.device).to(dtype)
+    else:
+        inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    ref_inp = to_reference(inp, False)
+
+    ref_out = torch.ops.aten.nonzero_numpy(ref_inp)
+    with flag_gems.use_gems():
+        res_out = torch.ops.aten.nonzero_numpy(inp)
+
+    assert len(res_out) == len(ref_out), "Number of output tensors should match"
+    for res_t, ref_t in zip(res_out, ref_out):
+        gems_assert_equal(res_t, ref_t)