Numba Unique: align with Python implementation

Author: ricardoV94

pytensor/link/numba/dispatch/extra_ops.py

@@ -13,6 +13,7 @@
     register_funcify_and_cache_key,
     register_funcify_default_op_cache_key,
 )
+from pytensor.npy_2_compat import old_np_unique
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.extra_ops import (
     Bartlett,
@@ -241,10 +242,17 @@ def unique(x):
         @numba_basic.numba_njit
         def unique(x):
             with numba.objmode(ret=ret_sig):
-                ret = np.unique(x, return_index, return_inverse, return_counts, axis)
+                ret = old_np_unique(
+                    x,
+                    return_index=return_index,
+                    return_inverse=return_inverse,
+                    return_counts=return_counts,
+                    axis=axis,
+                )
             return ret
 
-    return unique
+    cache_version = 1
+    return unique, cache_version
 
 
 @register_funcify_and_cache_key(UnravelIndex)

tests/link/numba/test_extra_ops.py

@@ -296,6 +296,14 @@ def test_Repeat(x, repeats, axis, exc):
             True,
             UserWarning,
         ),
+        (
+            (pt.lmatrix(), np.array([[1, 1], [1, 1], [2, 2]], dtype="int64")),
+            None,
+            True,
+            True,
+            True,
+            UserWarning,
+        ),
     ],
 )
 def test_Unique(x, axis, return_index, return_inverse, return_counts, exc):