Remove useless order kwarg from Sort and add numpy 2.0 stable kwarg

pytensor/link/jax/dispatch/sort.py

@@ -6,7 +6,9 @@ from pytensor.tensor.sort import SortOp
 
 @jax_funcify.register(SortOp)
 def jax_funcify_Sort(op, **kwargs):
+    stable = op.kind == "stable"
+
     def sort(arr, axis):
-        return jnp.sort(arr, axis=axis)
+        return jnp.sort(arr, axis=axis, stable=stable)
 
     return sort

pytensor/tensor/sort.py

+import typing
+
 import numpy as np
 
 from pytensor.gradient import grad_undefined
@@ -9,20 +11,34 @@ from pytensor.tensor.math import eq, ge
 from pytensor.tensor.type import TensorType
 
 
+KIND = typing.Literal["quicksort", "mergesort", "heapsort", "stable"]
+KIND_VALUES = typing.get_args(KIND)
+
+
+def _parse_sort_args(kind: KIND | None, order, stable: bool | None) -> KIND:
+    if order is not None:
+        raise ValueError("The order argument is not applicable to PyTensor graphs")
+    if stable is not None and kind is not None:
+        raise ValueError("kind and stable cannot be set at the same time")
+    if stable:
+        kind = "stable"
+    elif kind is None:
+        kind = "quicksort"
+    if kind not in KIND_VALUES:
+        raise ValueError(f"kind must be one of {KIND_VALUES}, got {kind}")
+    return kind
+
+
 class SortOp(Op):
     """
     This class is a wrapper for numpy sort function.
 
     """
 
-    __props__ = ("kind", "order")
+    __props__ = ("kind",)
 
-    def __init__(self, kind, order=None):
+    def __init__(self, kind: KIND):
         self.kind = kind
-        self.order = order
-
-    def __str__(self):
-        return self.__class__.__name__ + f"{{{self.kind}, {self.order}}}"
 
     def make_node(self, input, axis=-1):
         input = as_tensor_variable(input)
@@ -33,7 +49,7 @@ class SortOp(Op):
     def perform(self, node, inputs, output_storage):
         a, axis = inputs
         z = output_storage[0]
-        z[0] = np.sort(a, int(axis), self.kind, self.order)
+        z[0] = np.sort(a, int(axis), self.kind)
 
     def infer_shape(self, fgraph, node, inputs_shapes):
         assert node.inputs[0].ndim == node.outputs[0].ndim
@@ -75,9 +91,9 @@ class SortOp(Op):
 
         # The goal is to get gradient wrt input from gradient
         # wrt sort(input, axis)
-        idx = argsort(a, axis, kind=self.kind, order=self.order)
+        idx = argsort(a, axis, kind=self.kind)
         # rev_idx is the reverse of previous argsort operation
-        rev_idx = argsort(idx, axis, kind=self.kind, order=self.order)
+        rev_idx = argsort(idx, axis, kind=self.kind)
         indices = []
         axis_data = switch(ge(axis.data, 0), axis.data, a.ndim + axis.data)
         for i in range(a.ndim):
@@ -101,7 +117,9 @@ class SortOp(Op):
     """
 
 
-def sort(a, axis=-1, kind="quicksort", order=None):
+def sort(
+    a, axis=-1, kind: KIND | None = None, order=None, *, stable: bool | None = None
+):
     """
 
     Parameters
@@ -111,12 +129,12 @@ def sort(a, axis=-1, kind="quicksort", order=None):
     axis: TensorVariable
         Axis along which to sort. If None, the array is flattened before
         sorting.
-    kind: {'quicksort', 'mergesort', 'heapsort'}, optional
-        Sorting algorithm. Default is 'quicksort'.
+    kind: {'quicksort', 'mergesort', 'heapsort' 'stable'}, optional
+        Sorting algorithm. Default is 'quicksort' unless stable is defined.
     order: list, optional
-        When `a` is a structured array, this argument specifies which
-        fields to compare first, second, and so on. This list does not
-        need to include all of the fields.
+        For compatibility with numpy sort signature. Cannot be specified.
+    stable: bool, optional
+        Same as specifying kind = 'stable'. Cannot be specified at the same time as kind
 
     Returns
     -------
@@ -124,10 +142,12 @@ def sort(a, axis=-1, kind="quicksort", order=None):
         A sorted copy of an array.
 
     """
+    kind = _parse_sort_args(kind, order, stable)
+
     if axis is None:
         a = a.flatten()
         axis = 0
-    return SortOp(kind, order)(a, axis)
+    return SortOp(kind)(a, axis)
 
 
 class ArgSortOp(Op):
@@ -136,14 +156,10 @@ class ArgSortOp(Op):
 
     """
 
-    __props__ = ("kind", "order")
+    __props__ = ("kind",)
 
-    def __init__(self, kind, order=None):
+    def __init__(self, kind: KIND):
         self.kind = kind
-        self.order = order
-
-    def __str__(self):
-        return self.__class__.__name__ + f"{{{self.kind}, {self.order}}}"
 
     def make_node(self, input, axis=-1):
         input = as_tensor_variable(input)
@@ -158,7 +174,7 @@ class ArgSortOp(Op):
         a, axis = inputs
         z = output_storage[0]
         z[0] = _asarray(
-            np.argsort(a, int(axis), self.kind, self.order),
+            np.argsort(a, int(axis), self.kind),
             dtype=node.outputs[0].dtype,
         )
 
@@ -192,7 +208,9 @@ class ArgSortOp(Op):
     """
 
 
-def argsort(a, axis=-1, kind="quicksort", order=None):
+def argsort(
+    a, axis=-1, kind: KIND | None = None, order=None, stable: bool | None = None
+):
     """
     Returns the indices that would sort an array.
 
@@ -202,7 +220,8 @@ def argsort(a, axis=-1, kind="quicksort", order=None):
     order.
 
     """
+    kind = _parse_sort_args(kind, order, stable)
     if axis is None:
         a = a.flatten()
         axis = 0
-    return ArgSortOp(kind, order)(a, axis)
+    return ArgSortOp(kind)(a, axis)

tests/tensor/test_sort.py

 import numpy as np
+import pytest
 
 import pytensor
 from pytensor.tensor.sort import ArgSortOp, SortOp, argsort, sort
@@ -65,13 +66,12 @@ class TestSort:
             utt.assert_allclose(gv, gt)
 
     def test5(self):
-        a1 = SortOp("mergesort", [])
-        a2 = SortOp("quicksort", [])
+        a1 = SortOp("mergesort")
+        a2 = SortOp("quicksort")
 
-        # All the below should give true
         assert a1 != a2
-        assert a1 == SortOp("mergesort", [])
-        assert a2 == SortOp("quicksort", [])
+        assert a1 == SortOp("mergesort")
+        assert a2 == SortOp("quicksort")
 
     def test_None(self):
         a = dmatrix()
@@ -208,14 +208,11 @@ def test_argsort():
         utt.assert_allclose(gv, gt)
 
     # Example 5
-    a = dmatrix()
-    axis = lscalar()
-    a1 = ArgSortOp("mergesort", [])
-    a2 = ArgSortOp("quicksort", [])
-    # All the below should give true
+    a1 = ArgSortOp("mergesort")
+    a2 = ArgSortOp("quicksort")
     assert a1 != a2
-    assert a1 == ArgSortOp("mergesort", [])
-    assert a2 == ArgSortOp("quicksort", [])
+    assert a1 == ArgSortOp("mergesort")
+    assert a2 == ArgSortOp("quicksort")
 
     # Example 6: Testing axis=None
     a = dmatrix()
@@ -237,3 +234,22 @@ def test_argsort_grad():
 
     data = rng.random((2, 3, 3)).astype(pytensor.config.floatX)
     utt.verify_grad(lambda x: argsort(x, axis=2), [data])
+
+
+@pytest.mark.parametrize("func", (sort, argsort))
+def test_parse_sort_args(func):
+    x = matrix("x")
+
+    assert func(x).owner.op.kind == "quicksort"
+    assert func(x, stable=True).owner.op.kind == "stable"
+
+    with pytest.raises(ValueError, match="kind must be one of"):
+        func(x, kind="hanoi")
+
+    with pytest.raises(
+        ValueError, match="kind and stable cannot be set at the same time"
+    ):
+        func(x, kind="quicksort", stable=True)
+
+    with pytest.raises(ValueError, match="order argument is not applicable"):
+        func(x, order=[])