From 0fe14ad9cc2b7fa5b50dcf90234ba442bf48e00a Mon Sep 17 00:00:00 2001 From: guillaume-osmo Date: Wed, 26 Aug 2026 16:17:23 +0200 Subject: [PATCH] Pass pad_width as Python ints in the MLX Pad dispatch `mlx_funcify_pad` forwarded `pad_width` straight to `mx.pad`, but the linker typifies every input to `mx.array` while `mx.pad` takes an int or a list of (before, after) int pairs, so no pad worked on this backend -- both `test_mlx_pad[constant_default]` and `test_mlx_pad[edge]` were already red on main. Resolve a constant `pad_width` at funcify time, as the `Reshape` dispatch does, since `mx.compile` forbids reading a traced array; fall back to a runtime conversion otherwise, with an explicit `NotImplementedError` if the value turns out not to be readable. Same root cause as #2386 and the sort/argsort `axis` fix. --- pytensor/link/mlx/dispatch/pad.py | 50 +++++++++++++++++++++++++++++-- tests/link/mlx/test_pad.py | 15 ++++++++++ 2 files changed, 62 insertions(+), 3 deletions(-) diff --git a/pytensor/link/mlx/dispatch/pad.py b/pytensor/link/mlx/dispatch/pad.py index f29bba2129..1013c17c46 100644 --- a/pytensor/link/mlx/dispatch/pad.py +++ b/pytensor/link/mlx/dispatch/pad.py @@ -1,9 +1,41 @@ import mlx.core as mx +from pytensor.graph.basic import Constant from pytensor.link.mlx.dispatch.basic import mlx_funcify from pytensor.tensor.pad import Pad +PAD_WIDTH_NOT_COMPATIBLE = """MLX requires a concrete value for `pad_width`. + +The linker typifies every input to `mx.array`, but `mx.pad` takes an int or a +list of (before, after) int pairs, and `mx.compile` forbids reading a traced +array. Use a constant `pad_width`. +""" + + +def _resolve_pad_width(pad_width_input): + """Return `pad_width` as Python ints at funcify time, or None if not static.""" + if isinstance(pad_width_input, Constant): + value = pad_width_input.data.tolist() + # A 2-d (n, 2) spec has to be a list of pairs; `mx.pad` rejects a + # nested list of lists in some versions, and an int stays an int. + if isinstance(value, list) and value and isinstance(value[0], list): + return [tuple(pair) for pair in value] + return value + return None + + +def _pad_width_at_runtime(pad_width): + if isinstance(pad_width, mx.array): + try: + pad_width = pad_width.tolist() + except ValueError as exc: + raise NotImplementedError(PAD_WIDTH_NOT_COMPATIBLE) from exc + if isinstance(pad_width, list) and pad_width and isinstance(pad_width[0], list): + return [tuple(pair) for pair in pad_width] + return pad_width + + @mlx_funcify.register(Pad) def mlx_funcify_pad(op, node, **kwargs): pad_mode = op.pad_mode @@ -16,17 +48,29 @@ def mlx_funcify_pad(op, node, **kwargs): "not per-side tuples like NumPy/JAX." ) + static_pad_width = _resolve_pad_width(node.inputs[1]) + def constant_pad_fn(x, pad_width, constant_values): - return mx.pad( - x, pad_width, mode="constant", constant_values=constant_values + # `mx.pad` needs Python ints; the linker hands us an `mx.array`. + width = ( + static_pad_width + if static_pad_width is not None + else _pad_width_at_runtime(pad_width) ) + return mx.pad(x, width, mode="constant", constant_values=constant_values) return constant_pad_fn elif pad_mode == "edge": + static_pad_width = _resolve_pad_width(node.inputs[1]) def edge_pad_fn(x, pad_width): - return mx.pad(x, pad_width, mode="edge") + width = ( + static_pad_width + if static_pad_width is not None + else _pad_width_at_runtime(pad_width) + ) + return mx.pad(x, width, mode="edge") return edge_pad_fn diff --git a/tests/link/mlx/test_pad.py b/tests/link/mlx/test_pad.py index b7e8c96b81..51abd11191 100644 --- a/tests/link/mlx/test_pad.py +++ b/tests/link/mlx/test_pad.py @@ -49,3 +49,18 @@ def test_mlx_pad_non_scalar_constant_values(): NotImplementedError, match="only accepts a scalar constant_values" ): compare_mlx_and_py([x_pt], [res], [np.ones((3, 3))]) + + +@pytest.mark.parametrize("mode", ["constant", "edge"]) +@pytest.mark.parametrize( + "pad_width", [2, (1, 2), ((1, 2), (3, 0))], ids=["scalar", "pair", "per_axis"] +) +def test_mlx_pad_width_forms(mode, pad_width): + # `pad_width` reaches the dispatch as an `mx.array` because the linker + # typifies every input, while `mx.pad` takes an int or a list of int pairs, + # so nothing on this backend padded at all. Same root cause as #2386. + x = pt.tensor("x", shape=(3, 4), dtype="float32") + x_val = np.random.default_rng(0).normal(size=(3, 4)).astype("float32") + compare_mlx_and_py( + [x], [pt.pad(x, pad_width=pad_width, mode=mode)], [x_val], mlx_mode="MLX" + )