Modified the tests and the dispatch for efficiency

Author: Abhinav-Khot

pytensor/link/numba/dispatch/basic.py

@@ -749,18 +749,11 @@ def ifelse(cond, *args):
 
 @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)
+        if a.ndim == 1:
+            return result_tuple[0]
         return list(result_tuple)
 
     return nonzero

tests/link/numba/test_basic.py

@@ -31,7 +31,6 @@
 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
 
@@ -288,7 +287,6 @@ def assert_fn(x, y):
     )
     test_inputs_copy = (inp.copy() for inp in test_inputs) if inplace else test_inputs
     numba_res = pytensor_numba_fn(*test_inputs_copy)
-
     if isinstance(graph_outputs, tuple | list):
         for j, p in zip(numba_res, py_res, strict=True):
             assert_fn(j, p)
@@ -901,17 +899,9 @@ def test_function_overhead(mode, benchmark):
     [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()
+    a = pt.tensor("a", shape=(None,) * input_data.ndim)
 
-    node = nonzero_op.make_node(a)
-    graph_outputs = node.outputs
+    graph_outputs = pt.nonzero(a)
 
     compare_numba_and_py(
         graph_inputs=[a], graph_outputs=graph_outputs, test_inputs=[input_data]