added numba backend and testsfor Nonzero

Author: Abhinav-Khot

pytensor/link/numba/dispatch/basic.py

@@ -33,6 +33,7 @@
 from pytensor.scalar.basic import ScalarType
 from pytensor.scalar.math import Softplus
 from pytensor.sparse import SparseTensorType
+from pytensor.tensor.basic import Nonzero
 from pytensor.tensor.blas import BatchedDot
 from pytensor.tensor.math import Dot
 from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
@@ -744,3 +745,22 @@ def ifelse(cond, *args):
             return res[0]
 
     return ifelse
+
+
+@numba_funcify.register(Nonzero)
+def numba_funcify_Nonzero(op, node, **kwargs):
+    a = node.inputs[0]
+
+    if a.ndim == 0:
+        raise ValueError("Nonzero only supports non-scalar arrays.")
+
+    @numba_njit
+    def nonzero(a):
+        if a.ndim == 1:
+            indices = np.where(a != 0)[0]
+            return indices.astype(np.int64)
+
+        result_tuple = np.nonzero(a)
+        return list(result_tuple)
+
+    return nonzero

tests/link/numba/test_basic.py

@@ -31,6 +31,7 @@
 from pytensor.raise_op import assert_op
 from pytensor.scalar.basic import ScalarOp, as_scalar
 from pytensor.tensor import blas
+from pytensor.tensor.basic import Nonzero
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
 
@@ -893,3 +894,25 @@ def test_function_overhead(mode, benchmark):
     assert np.sum(fn(test_x)) == 1000
 
     benchmark(fn, test_x)
+
+
+@pytest.mark.parametrize(
+    "input_data",
+    [np.array([1, 0, 3]), np.array([[0, 1], [2, 0]]), np.array([[0, 0], [0, 0]])],
+)
+def test_Nonzero(input_data):
+    if input_data.ndim == 0:
+        a = pt.scalar("a")
+    elif input_data.ndim == 1:
+        a = pt.vector("a")
+    elif input_data.ndim == 2:
+        a = pt.matrix("a")
+
+    nonzero_op = Nonzero()
+
+    node = nonzero_op.make_node(a)
+    graph_outputs = node.outputs
+
+    compare_numba_and_py(
+        graph_inputs=[a], graph_outputs=graph_outputs, test_inputs=[input_data]
+    )