From 8b5d226d63f3784645bb8d539a4ddb4f63b97c7b Mon Sep 17 00:00:00 2001 From: guillaume-osmo Date: Wed, 26 Aug 2026 15:49:37 +0200 Subject: [PATCH] Fix MLX Reshape: resolve the shape at funcify time (#2386) `mlx_funcify_Reshape` forwarded the shape input straight to `mx.reshape`, but the linker typifies every input to `mx.array` while `mx.reshape` only accepts a Python sequence of ints, so every reshape raised `TypeError`. Reading the array back at runtime is not an option either: the linker enables `mx.compile` by default and MLX forbids evaluating a traced array, so the shape has to be resolved when the dispatch is built. Prefer the statically inferred output shape (which already resolves any `-1`, and covers shapes read off another input such as `x.reshape(y.shape)`), then fall back to a constant shape input. A genuinely data-dependent shape now raises `NotImplementedError` with an explanation instead of an eval error from deep inside MLX. This turns four already-red tests in `tests/link/mlx/test_shape.py` green and adds regression tests under the full "MLX" mode. --- pytensor/link/mlx/dispatch/shape.py | 48 +++++++++++++++++++++++++++-- tests/link/mlx/test_shape.py | 26 ++++++++++++++++ 2 files changed, 71 insertions(+), 3 deletions(-) diff --git a/pytensor/link/mlx/dispatch/shape.py b/pytensor/link/mlx/dispatch/shape.py index ab5205d7b4..af7a1ce1c3 100644 --- a/pytensor/link/mlx/dispatch/shape.py +++ b/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 diff --git a/tests/link/mlx/test_shape.py b/tests/link/mlx/test_shape.py index c19247ad0e..b7d33dba6d 100644 --- a/tests/link/mlx/test_shape.py +++ b/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", + )