Add nonzero_numpy operator implementation, tests and benchmark

Author: Schopenhauer-loves-Hegel

benchmark/test_reduction_perf.py

@@ -159,6 +159,13 @@ def mse_loss_input_fn(shape, cur_dtype, device):
             FLOAT_DTYPES + INT_DTYPES + BOOL_DTYPES,
             marks=pytest.mark.nonzero,
         ),
+        pytest.param(
+            "nonzero_numpy",
+            torch.ops.aten.nonzero_numpy,
+            unary_input_fn,
+            FLOAT_DTYPES + INT_DTYPES + BOOL_DTYPES,
+            marks=pytest.mark.nonzero_numpy,
+        ),
         pytest.param(
             "cross_entropy_loss",
             torch.nn.functional.cross_entropy,

src/flag_gems/__init__.py

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

src/flag_gems/ops/__init__.py

@@ -152,6 +152,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,
@@ -435,6 +436,7 @@
     "nll_loss2d_backward",
     "nll_loss2d_forward",
     "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_reduction_ops.py

@@ -513,6 +513,29 @@ def test_accuracy_nonzero(shape, dtype):
     gems_assert_equal(res_out, ref_out)
 
 
+@pytest.mark.nonzero_numpy
+@pytest.mark.parametrize("shape", NONZERO_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES + INT_DTYPES + [torch.bool])
+def test_accuracy_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)
+
+
 @pytest.mark.count_nonzero
 @pytest.mark.parametrize("shape", REDUCTION_SHAPES)
 @pytest.mark.parametrize("dtype", FLOAT_DTYPES + INT_DTYPES + [torch.bool])