Move JAX sort dispatch to its own module

Author: ricardoV94

pytensor/link/jax/dispatch/__init__.py

@@ -14,5 +14,6 @@
 import pytensor.link.jax.dispatch.scan
 import pytensor.link.jax.dispatch.sparse
 import pytensor.link.jax.dispatch.blockwise
+import pytensor.link.jax.dispatch.sort
 
 # isort: on

pytensor/link/jax/dispatch/sort.py

@@ -0,0 +1,12 @@
+from jax import numpy as jnp
+
+from pytensor.link.jax.dispatch import jax_funcify
+from pytensor.tensor.sort import SortOp
+
+
+@jax_funcify.register(SortOp)
+def jax_funcify_Sort(op, **kwargs):
+    def sort(arr, axis):
+        return jnp.sort(arr, axis=axis)
+
+    return sort

pytensor/link/jax/dispatch/tensor_basic.py

@@ -22,7 +22,6 @@
 )
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.shape import Shape_i
-from pytensor.tensor.sort import SortOp
 
 
 ARANGE_CONCRETE_VALUE_ERROR = """JAX requires the arguments of `jax.numpy.arange`
@@ -206,11 +205,3 @@ def tri(*args):
         return jnp.tri(*args, dtype=op.dtype)
 
     return tri
-
-
-@jax_funcify.register(SortOp)
-def jax_funcify_Sort(op, **kwargs):
-    def sort(arr, axis):
-        return jnp.sort(arr, axis=axis)
-
-    return sort

tests/link/jax/__init__.py

@@ -0,0 +1,16 @@
+import numpy as np
+import pytest
+
+from pytensor.graph import FunctionGraph
+from pytensor.tensor import matrix
+from pytensor.tensor.sort import sort
+from tests.link.jax.test_basic import compare_jax_and_py
+
+
+@pytest.mark.parametrize("axis", [None, -1])
+def test_sort(axis):
+    x = matrix("x", shape=(2, 2), dtype="float64")
+    out = sort(x, axis=axis)
+    fgraph = FunctionGraph([x], [out])
+    arr = np.array([[1.0, 4.0], [5.0, 2.0]])
+    compare_jax_and_py(fgraph, [arr])

tests/link/jax/test_sort.py

@@ -218,15 +218,6 @@ def test_tri():
     compare_jax_and_py(fgraph, [])
 
 
-@pytest.mark.parametrize("axis", [None, -1])
-def test_sort(axis):
-    x = matrix("x", shape=(2, 2), dtype="float64")
-    out = pytensor.tensor.sort(x, axis=axis)
-    fgraph = FunctionGraph([x], [out])
-    arr = np.array([[1.0, 4.0], [5.0, 2.0]])
-    compare_jax_and_py(fgraph, [arr])
-
-
 def test_tri_nonconcrete():
     """JAX cannot JIT-compile `jax.numpy.tri` when arguments are not concrete values."""