Implemented Sort/Argsort Ops in PyTorch

Author: twaclaw

pytensor/link/pytorch/dispatch/__init__.py

@@ -5,4 +5,5 @@
 import pytensor.link.pytorch.dispatch.scalar
 import pytensor.link.pytorch.dispatch.elemwise
 import pytensor.link.pytorch.dispatch.extra_ops
+import pytensor.link.pytorch.dispatch.sort
 # isort: on

pytensor/link/pytorch/dispatch/sort.py

@@ -0,0 +1,25 @@
+import torch
+
+from pytensor.link.pytorch.dispatch.basic import pytorch_funcify
+from pytensor.tensor.sort import ArgSortOp, SortOp
+
+
+@pytorch_funcify.register(SortOp)
+def pytorch_funcify_Sort(op, **kwargs):
+    stable = op.kind == "stable"
+
+    def sort(arr, axis):
+        sorted, _ = torch.sort(arr, dim=axis, stable=stable)
+        return sorted
+
+    return sort
+
+
+@pytorch_funcify.register(ArgSortOp)
+def pytorch_funcify_ArgSort(op, **kwargs):
+    stable = op.kind == "stable"
+
+    def argsort(arr, axis):
+        return torch.argsort(arr, dim=axis, stable=stable)
+
+    return argsort

tests/link/pytorch/test_sort.py

@@ -0,0 +1,23 @@
+import numpy as np
+import pytest
+
+from pytensor.graph import FunctionGraph
+from pytensor.tensor import matrix
+from pytensor.tensor.sort import argsort, sort
+from tests.link.pytorch.test_basic import compare_pytorch_and_py
+
+
+@pytest.mark.parametrize("axis", [0, 1, None])
+@pytest.mark.parametrize("func", (sort, argsort))
+def test_sort(func, axis):
+    x = matrix("x", shape=(2, 2), dtype="float64")
+    out = func(x, axis=axis)
+    fgraph = FunctionGraph([x], [out])
+    arr = np.array([[1.0, 4.0], [5.0, 2.0]])
+
+    # TODO: remove condition once Reshape is implemented
+    if axis is None:
+        with pytest.raises(NotImplementedError):
+            compare_pytorch_and_py(fgraph, [arr])
+    else:
+        compare_pytorch_and_py(fgraph, [arr])