default to supported kind and add warning

Signed-off-by: Victor Garcia Reolid <[email protected]>

Author: Victor Garcia Reolid

pytensor/link/numba/dispatch/basic.py

@@ -435,18 +435,18 @@ def shape_i(x):
 
 @numba_funcify.register(SortOp)
 def numba_funcify_SortOp(op, node, **kwargs):
-    if op.kind == "quicksort":
-
-        @numba_njit
-        def sort_f(a, axis):
-            return np.sort(a)  # numba supports sort without arguments
-    else:
-        ret_sig = get_numba_type(node.outputs[0].type)
+    @numba_njit
+    def sort_f(a, axis):
+        return np.sort(a)  # numba supports sort without arguments
 
-        def sort_f(a, axis):
-            with numba.objmode(ret=ret_sig):
-                ret = np.sort(a, axis=axis, kind=op.kind)
-            return ret
+    if op.kind != "quicksort":
+        warnings.warn(
+            (
+                f'Numba function sort doesn\'t support kind="{op.kind}"'
+                " switching to `quicksort`."
+            ),
+            UserWarning,
+        )
 
     return sort_f
 
@@ -460,17 +460,20 @@ def argsort_f(a, axis):
 
         return argsort_f
 
-    if op.kind in ["quicksort", "mergesort"]:
-        return argsort_f_kind(op.kind)
-    else:
-        ret_sig = get_numba_type(node.outputs[0].type)
+    kind = op.kind
 
-        def argsort_f(a, axis):
-            with numba.objmode(ret=ret_sig):
-                ret = np.argsort(a, axis=axis, kind=op.kind)
-            return ret
+    if kind in ["quicksort", "mergesort"]:
+        return argsort_f_kind(kind)
+    else:
+        warnings.warn(
+            (
+                f'Numba function argsort doesn\'t support kind="{op.kind}"'
+                " switching to `quicksort`."
+            ),
+            UserWarning,
+        )
 
-    return argsort_f
+        return argsort_f_kind("quicksort")
 
 
 @numba.extending.intrinsic

tests/link/numba/test_basic.py

@@ -379,18 +379,47 @@ def test_Shape(x, i):
     compare_numba_and_py([], [g], [])
 
 
-@pytest.mark.parametrize("kind", ["quicksort"])
-def test_Sort(kind):
+@pytest.mark.parametrize(
+    "kind, exc",
+    [
+        ["quicksort", None],
+        ["mergesort", UserWarning],
+        ["heapsort", UserWarning],
+        ["stable", UserWarning],
+    ],
+)
+def test_Sort(kind, exc):
     x = [5, 4, 3, 2, 1]
+
     g = SortOp(kind)(pt.as_tensor_variable(x))
+
+    if exc:
+        with pytest.warns(exc):
+            compare_numba_and_py([], [g], [])
+    else:
+        compare_numba_and_py([], [g], [])
+
     compare_numba_and_py([], [g], [])
 
 
-@pytest.mark.parametrize("kind", ["quicksort", "mergesort"])
-def test_ArgSort(kind):
+@pytest.mark.parametrize(
+    "kind, exc",
+    [
+        ["quicksort", None],
+        ["mergesort", None],
+        ["heapsort", UserWarning],
+        ["stable", UserWarning],
+    ],
+)
+def test_ArgSort(kind, exc):
     x = [5, 4, 3, 2, 1]
     g = ArgSortOp(kind)(pt.as_tensor_variable(x))
-    compare_numba_and_py([], [g], [])
+
+    if exc:
+        with pytest.warns(exc):
+            compare_numba_and_py([], [g], [])
+    else:
+        compare_numba_and_py([], [g], [])
 
 
 @pytest.mark.parametrize(