Pass pad_width as Python ints in the MLX Pad dispatch

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
 

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