feat: support numba compiled sort and argsort functions

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

Author: Victor Garcia Reolid

pytensor/link/numba/dispatch/basic.py

@@ -37,6 +37,7 @@
 from pytensor.tensor.math import Dot
 from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
 from pytensor.tensor.slinalg import Solve
+from pytensor.tensor.sort import ArgSortOp, SortOp
 from pytensor.tensor.type import TensorType
 from pytensor.tensor.type_other import MakeSlice, NoneConst
 
@@ -432,6 +433,46 @@ def shape_i(x):
     return shape_i
 
 
+@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)
+
+        def sort_f(a, axis):
+            with numba.objmode(ret=ret_sig):
+                ret = np.sort(a, axis=axis, kind=op.kind)
+            return ret
+
+    return sort_f
+
+
+@numba_funcify.register(ArgSortOp)
+def numba_funcify_ArgSortOp(op, node, **kwargs):
+    def argsort_f_kind(kind):
+        @numba_njit
+        def argsort_f(a, axis):
+            return np.argsort(a, kind=kind)
+
+        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)
+
+        def argsort_f(a, axis):
+            with numba.objmode(ret=ret_sig):
+                ret = np.argsort(a, axis=axis, kind=op.kind)
+            return ret
+
+    return argsort_f
+
+
 @numba.extending.intrinsic
 def direct_cast(typingctx, val, typ):
     if isinstance(typ, numba.types.TypeRef):

tests/link/numba/test_basic.py

@@ -33,6 +33,7 @@
 from pytensor.tensor import blas
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
+from pytensor.tensor.sort import ArgSortOp, SortOp
 
 
 if TYPE_CHECKING:
@@ -378,6 +379,20 @@ def test_Shape(x, i):
     compare_numba_and_py([], [g], [])
 
 
+@pytest.mark.parametrize("kind", ["quicksort"])
+def test_Sort(kind):
+    x = [5, 4, 3, 2, 1]
+    g = SortOp(kind)(pt.as_tensor_variable(x))
+    compare_numba_and_py([], [g], [])
+
+
+@pytest.mark.parametrize("kind", ["quicksort", "mergesort"])
+def test_ArgSort(kind):
+    x = [5, 4, 3, 2, 1]
+    g = ArgSortOp(kind)(pt.as_tensor_variable(x))
+    compare_numba_and_py([], [g], [])
+
+
 @pytest.mark.parametrize(
     "v, shape, ndim",
     [