From 8b5d226d63f3784645bb8d539a4ddb4f63b97c7b Mon Sep 17 00:00:00 2001 From: guillaume-osmo Date: Wed, 26 Aug 2026 15:49:37 +0200 Subject: [PATCH 1/2] 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", + ) From 11498e780c2d9d3e463b55cbdc2bd699e90fc401 Mon Sep 17 00:00:00 2001 From: guillaume-osmo Date: Wed, 26 Aug 2026 16:06:04 +0200 Subject: [PATCH 2/2] Pass axis as a Python int in the MLX sort/argsort dispatches `mlx_funcify_Sort` and `mlx_funcify_ArgSort` forwarded the `axis` input straight to `mx.sort` / `mx.argsort`, but the linker typifies every input to `mx.array` and both functions require a Python int, so nothing on this backend sorted at all -- every call raised `TypeError`. Same root cause as the `Reshape` dispatch (#2386), which `axis=None` also goes through, since it flattens first. All five tests in `tests/link/mlx/test_sort.py` were already red on main; they pass now, and a parametrized test covers positive, negative and `None` axes on a 3-d input. Found while fixing #2385: the gradient of a batched `solve` reaches `argsort` through the pivot-to-permutation step. --- pytensor/link/mlx/dispatch/sort.py | 8 ++++++-- tests/link/mlx/test_sort.py | 14 +++++++++++++- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/pytensor/link/mlx/dispatch/sort.py b/pytensor/link/mlx/dispatch/sort.py index 61e13dbba6..9afcf8326f 100644 --- a/pytensor/link/mlx/dispatch/sort.py +++ b/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 diff --git a/tests/link/mlx/test_sort.py b/tests/link/mlx/test_sort.py index 8d37a76ab1..cde22a8c5c 100644 --- a/tests/link/mlx/test_sort.py +++ b/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")