Fix MLX Reshape: resolve the shape at funcify time

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

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",
+    )