Pass axis as a Python int in the MLX sort/argsort dispatches

pytensor/link/mlx/dispatch/shape.py

@@ -1,5 +1,6 @@
 import mlx.core as mx
 
+from pytensor.graph.basic import Constant
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
 from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
 
@@ -36,9 +37,50 @@ def shape_i(x):
     return shape_i
 
 
+SHAPE_NOT_COMPATIBLE = """MLX requires a concrete value for the `shape` argument of `mx.reshape`.
+
+The linker typifies every input to `mx.array`, and `mx.compile` traces the graph
+with static shapes, so a shape whose *values* are only known at runtime cannot be
+read back. Use a constant shape, or one that PyTensor's shape inference can
+resolve statically:
+
+>>> import pytensor.tensor as pt
+>>> x = pt.ones((6, 4))
+>>> y = x.reshape((24,))          # constant
+>>> mat = pt.matrix("mat", shape=(6, 4))
+>>> y = mat.reshape(mat.shape)    # statically resolvable
+"""
+
+
 @mlx_funcify.register(Reshape)
-def mlx_funcify_Reshape(op, **kwargs):
-    def reshape(x, shp):
-        return mx.reshape(x, shp)
+def mlx_funcify_Reshape(op, node, **kwargs):
+    # `mx.reshape` wants a Python sequence of ints, but the linker typifies the
+    # shape input to `mx.array` and `mx.compile` forbids reading a traced array,
+    # so the shape has to be resolved at funcify time (#2386).
+    static_shape = node.outputs[0].type.shape
+    shape_input = node.inputs[1]
+
+    if not any(dim is None for dim in static_shape):
+        # Shape inference already resolved every dimension, including any -1.
+        target = tuple(static_shape)
+    elif isinstance(shape_input, Constant):
+        target = tuple(int(dim) for dim in shape_input.data)
+    else:
+        target = None
+
+    if target is not None:
+
+        def reshape(x, shp):
+            return mx.reshape(x, target)
+
+    else:
+
+        def reshape(x, shp):
+            if isinstance(shp, mx.array):
+                try:
+                    shp = shp.tolist()
+                except ValueError as exc:
+                    raise NotImplementedError(SHAPE_NOT_COMPATIBLE) from exc
+            return mx.reshape(x, tuple(shp))
 
     return reshape

pytensor/link/mlx/dispatch/sort.py

@@ -17,7 +17,10 @@ def mlx_funcify_Sort(op, **kwargs):
         )
 
     def sort(x, axis):
-        return mx.sort(x, axis=axis)
+        # The linker typifies every input to `mx.array`, but `mx.sort` wants a
+        # Python int for `axis`, so no sort worked on this backend at all.
+        # Same root cause as the `Reshape` dispatch (#2386).
+        return mx.sort(x, axis=int(axis))
 
     return sort
 
@@ -33,6 +36,7 @@ def mlx_funcify_ArgSort(op, **kwargs):
         )
 
     def argsort(x, axis):
-        return mx.argsort(x, axis=axis)
+        # As for `sort`: `axis` arrives as an `mx.array` and must be an int.
+        return mx.argsort(x, axis=int(axis))
 
     return argsort

tests/link/mlx/test_shape.py

@@ -117,3 +117,29 @@ def test_mlx_compile_ops():
     x = ViewOp()(pt.as_tensor_variable(x_np))
 
     compare_mlx_and_py([], [x], [])
+
+
+def test_mlx_Reshape_full_mlx_mode():
+    # Under the full "MLX" mode the linker typifies the shape input to an
+    # ``mx.array``, which ``mx.reshape`` rejects outright, so every reshape
+    # raised ``TypeError`` (#2386). The shape has to be resolved at funcify
+    # time instead.
+    x = pt.matrix("x", shape=(6, 4), dtype="float32")
+    x_val = np.arange(24, dtype="float32").reshape(6, 4)
+
+    for shape in ((24,), (4, 6), (2, 12), (-1, 3), (3, -1), (2, 3, 4)):
+        compare_mlx_and_py([x], [reshape(x, shape)], [x_val], mlx_mode="MLX")
+
+
+def test_mlx_Reshape_shape_from_other_input():
+    # A shape read off another input is not a ``Constant``, but PyTensor's shape
+    # inference still resolves it statically, so it must compile under "MLX".
+    x = pt.matrix("x", shape=(6, 4), dtype="float32")
+    y = vector("y", shape=(24,), dtype="float32")
+
+    compare_mlx_and_py(
+        [x, y],
+        [reshape(x, y.shape)],
+        [np.arange(24, dtype="float32").reshape(6, 4), np.zeros(24, dtype="float32")],
+        mlx_mode="MLX",
+    )

tests/link/mlx/test_sort.py

@@ -2,7 +2,7 @@
 import pytest
 
 from pytensor.tensor.sort import argsort, sort
-from pytensor.tensor.type import matrix
+from pytensor.tensor.type import matrix, tensor
 from tests.link.mlx.test_basic import compare_mlx_and_py
 
 
@@ -20,3 +20,15 @@ def test_sort_invalid_kind_warning():
     z = sort(x, axis=-1, kind="mergesort")
     with pytest.warns(UserWarning, match="MLX sort does not support the kind argument"):
         z.eval({x: np.array([[3.0, 1.0], [2.0, 4.0]])}, mode="MLX")
+
+
+@pytest.mark.parametrize("axis", [None, 0, 1, -1, -2])
+@pytest.mark.parametrize("func", (sort, argsort))
+def test_sort_axis_variants(func, axis):
+    # `axis` reaches the dispatch as an `mx.array` because the linker typifies
+    # every input, and `mx.sort`/`mx.argsort` require a Python int, so nothing
+    # on this backend sorted at all. `axis=None` additionally flattens through
+    # `Reshape`, which had the same root cause (#2386).
+    x = tensor("x", shape=(2, 3, 4), dtype="float32")
+    arr = np.random.default_rng(0).normal(size=(2, 3, 4)).astype("float32")
+    compare_mlx_and_py([x], [func(x, axis=axis)], [arr], mlx_mode="MLX")