Cleanup Op dispatchers

Author: ricardoV94

pytensor/link/numba/dispatch/sort.py

@@ -9,6 +9,15 @@
 
 @numba_funcify.register(SortOp)
 def numba_funcify_SortOp(op, node, **kwargs):
+    if op.kind != "quicksort":
+        warnings.warn(
+            (
+                f'Numba function sort doesn\'t support kind="{op.kind}"'
+                " switching to `quicksort`."
+            ),
+            UserWarning,
+        )
+
     @numba_njit
     def sort_f(a, axis):
         axis = axis.item()
@@ -19,41 +28,11 @@ def sort_f(a, axis):
 
         return a_sorted_swapped
 
-    if op.kind != "quicksort":
-        warnings.warn(
-            (
-                f'Numba function sort doesn\'t support kind="{op.kind}"'
-                " switching to `quicksort`."
-            ),
-            UserWarning,
-        )
-
     return sort_f
 
 
 @numba_funcify.register(ArgSortOp)
 def numba_funcify_ArgSortOp(op, node, **kwargs):
-    def argsort_f_kind(kind):
-        @numba_njit
-        def argort_vec(X, axis):
-            axis = axis.item()
-
-            Y = np.swapaxes(X, axis, 0)
-            result = np.empty_like(Y, dtype="int64")
-
-            indices = list(np.ndindex(Y.shape[1:]))
-
-            for idx in indices:
-                result[(slice(None), *idx)] = np.argsort(
-                    Y[(slice(None), *idx)], kind=kind
-                )
-
-            result = np.swapaxes(result, 0, axis)
-
-            return result
-
-        return argort_vec
-
     kind = op.kind
 
     if kind not in ["quicksort", "mergesort"]:
@@ -66,4 +45,19 @@ def argort_vec(X, axis):
             UserWarning,
         )
 
-    return argsort_f_kind(kind)
+    @numba_njit
+    def argort_f(X, axis):
+        axis = axis.item()
+
+        Y = np.swapaxes(X, axis, 0)
+        result = np.empty_like(Y, dtype="int64")
+
+        indices = list(np.ndindex(Y.shape[1:]))
+
+        for idx in indices:
+            result[(slice(None), *idx)] = np.argsort(Y[(slice(None), *idx)], kind=kind)
+
+        result = np.swapaxes(result, 0, axis)
+        return result
+
+    return argort_f