feat: support axis

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

Author: Victor Garcia Reolid

pytensor/link/numba/dispatch/basic.py

@@ -11,7 +11,7 @@
 import scipy
 import scipy.special
 from llvmlite import ir
-from numba import types
+from numba import prange, types
 from numba.core.errors import NumbaWarning, TypingError
 from numba.cpython.unsafe.tuple import tuple_setitem  # noqa: F401
 from numba.extending import box, overload
@@ -437,7 +437,14 @@ def shape_i(x):
 def numba_funcify_SortOp(op, node, **kwargs):
     @numba_njit
     def sort_f(a, axis):
-        return np.sort(a)  # numba supports sort without arguments
+        if not isinstance(axis, int):
+            axis = -1
+
+        a_swapped = np.swapaxes(a, axis, -1)
+        a_sorted = np.sort(a_swapped)
+        a_sorted_swapped = np.swapaxes(a_sorted, -1, axis)
+
+        return a_sorted_swapped
 
     if op.kind != "quicksort":
         warnings.warn(
@@ -455,10 +462,27 @@ def sort_f(a, axis):
 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)
+        def argort_vec(X, axis):
+            if axis > len(X.shape):
+                raise ValueError("Wrong axis.")
+
+            axis = axis.item()
+
+            Y = np.swapaxes(X, axis, 0)
+            result = np.empty_like(Y)
+
+            N = int(np.prod(np.array(Y.shape)[1:]))
+            indices = list(np.ndindex(Y.shape[1:]))
+
+            for i in prange(N):
+                idx = indices[i]
+                result[:, *idx] = np.argsort(Y[:, *idx], kind=kind)
+
+            result = np.swapaxes(result, 0, axis)
+
+            return result
 
-        return argsort_f
+        return argort_vec
 
     kind = op.kind
 

tests/link/numba/test_basic.py

@@ -380,45 +380,59 @@ def test_Shape(x, i):
 
 
 @pytest.mark.parametrize(
-    "kind, exc",
+    "x, axis, kind, exc",
     [
-        ["quicksort", None],
-        ["mergesort", UserWarning],
-        ["heapsort", UserWarning],
-        ["stable", UserWarning],
+        [[3, 2, 1], None, "quicksort", None],
+        [[], None, "quicksort", None],
+        [[[3, 2, 1], [5, 6, 7]], None, "quicksort", None],
+        [[3, 2, 1], None, "mergesort", UserWarning],
+        [[3, 2, 1], None, "heapsort", UserWarning],
+        [[3, 2, 1], None, "stable", UserWarning],
+        [[[3, 2, 1], [5, 6, 7]], 0, "quicksort", None],
+        [[[3, 2, 1], [5, 6, 7]], 1, "quicksort", None],
+        [[[3, 2, 1], [5, 6, 7]], -1, "quicksort", None],
+        [[3, 2, 1], 0, "quicksort", None],
+        [np.random.randint(0, 100, (40, 40, 40, 40)), 3, "quicksort", None],
     ],
 )
-def test_Sort(kind, exc):
-    x = [5, 4, 3, 2, 1]
+def test_Sort(x, axis, kind, exc):
+    if axis:
+        g = SortOp(kind)(pt.as_tensor_variable(x), axis)
+    else:
+        g = SortOp(kind)(pt.as_tensor_variable(x))
 
-    g = SortOp(kind)(pt.as_tensor_variable(x))
+    cm = contextlib.suppress() if not exc else pytest.warns(exc)
 
-    if exc:
-        with pytest.warns(exc):
-            compare_numba_and_py([], [g], [])
-    else:
+    with cm:
         compare_numba_and_py([], [g], [])
 
-    compare_numba_and_py([], [g], [])
-
 
 @pytest.mark.parametrize(
-    "kind, exc",
+    "x, axis, kind, exc",
     [
-        ["quicksort", None],
-        ["mergesort", None],
-        ["heapsort", UserWarning],
-        ["stable", UserWarning],
+        [[3, 2, 1], None, "quicksort", None],
+        [[], None, "quicksort", None],
+        [[[3, 2, 1], [5, 6, 7]], None, "quicksort", None],
+        [[3, 2, 1], None, "heapsort", UserWarning],
+        [[3, 2, 1], None, "stable", UserWarning],
+        [[[3, 2, 1], [5, 6, 7]], 0, "quicksort", None],
+        [[[3, 2, 1], [5, 6, 7]], None, "quicksort", None],
+        [[[3, 2, 1], [5, 6, 7]], 1, "quicksort", None],
+        [[[3, 2, 1], [5, 6, 7]], -1, "quicksort", None],
+        [[3, 2, 1], 0, "quicksort", None],
+        [np.random.randint(0, 10, (3, 2, 3)), 1, "quicksort", None],
+        [np.random.randint(0, 10, (3, 2, 3, 4, 4)), 2, "quicksort", None],
     ],
 )
-def test_ArgSort(kind, exc):
-    x = [5, 4, 3, 2, 1]
-    g = ArgSortOp(kind)(pt.as_tensor_variable(x))
-
-    if exc:
-        with pytest.warns(exc):
-            compare_numba_and_py([], [g], [])
+def test_ArgSort(x, axis, kind, exc):
+    if axis:
+        g = ArgSortOp(kind)(pt.as_tensor_variable(x), axis)
     else:
+        g = ArgSortOp(kind)(pt.as_tensor_variable(x))
+
+    cm = contextlib.suppress() if not exc else pytest.warns(exc)
+
+    with cm:
         compare_numba_and_py([], [g], [])