First draft JAX overload

Author: jessegrabowski

pytensor/link/jax/dispatch/__init__.py

@@ -7,6 +7,7 @@
 import pytensor.link.jax.dispatch.subtensor
 import pytensor.link.jax.dispatch.shape
 import pytensor.link.jax.dispatch.extra_ops
+import pytensor.link.jax.dispatch.pad
 import pytensor.link.jax.dispatch.nlinalg
 import pytensor.link.jax.dispatch.slinalg
 import pytensor.link.jax.dispatch.random

pytensor/link/jax/dispatch/pad.py

@@ -1,16 +1,20 @@
 import jax.numpy as jnp
 
 from pytensor.link.jax.dispatch import jax_funcify
-from pytensor.tensor.pad import Pad, allowed_kwargs
+from pytensor.tensor.pad import Pad
 
 
-@jax_funcify([Pad])
+fixed_kwargs = {"reflect": ["reflect_type"], "symmetric": ["reflect_type"]}
+
+
+@jax_funcify.register(Pad)
 def jax_funcify_pad(op, **kwargs):
     pad_mode = op.pad_mode
-    expected_kwargs = allowed_kwargs[pad_mode]
+    expected_kwargs = fixed_kwargs.get(pad_mode, {})
     mode_kwargs = {kwarg: getattr(op, kwarg) for kwarg in expected_kwargs}
 
-    def pad(x, pad_width, pad_mode=pad_mode):
-        return jnp.pad(x, pad_width=pad_width, pad_mode=pad_mode, **mode_kwargs)
+    def pad(x, pad_width, *args):
+        print(args)
+        return jnp.pad(x, pad_width, mode=pad_mode, **mode_kwargs)
 
     return pad

tests/link/jax/test_pad.py

@@ -0,0 +1,52 @@
+import numpy as np
+import pytest
+
+import pytensor.tensor as pt
+from pytensor import config
+from pytensor.graph import FunctionGraph
+from pytensor.tensor.pad import PadMode
+from tests.link.jax.test_basic import compare_jax_and_py
+
+
+jax = pytest.importorskip("jax")
+floatX = config.floatX
+RTOL = ATOL = 1e-6 if floatX.endswith("64") else 1e-3
+
+test_kwargs = {
+    "constant": {"constant_values": 0},
+    "linear_ramp": {"end_values": 0},
+    "maximum": {"stat_length": None},
+    "mean": {"stat_length": [[1, 2], [3, 3]]},
+    "median": {"stat_length": 2},
+    "reflect": {"reflect_type": "even"},
+    "symmetric": {"reflect_type": "even"},
+}
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [
+        "constant",
+        "edge",
+        "linear_ramp",
+        "wrap",
+        "symmetric",
+        "mean",
+        "maximum",
+        "minimum",
+    ],
+)
+def test_jax_pad(mode: PadMode):
+    x_pt = pt.dmatrix("x")
+    x = np.random.normal(size=(3, 3))
+    kwargs = test_kwargs.get(mode, {})
+
+    res = pt.pad(x_pt, mode=mode, pad_width=3, **kwargs)
+    res_fg = FunctionGraph([x_pt], [res])
+
+    compare_jax_and_py(
+        res_fg,
+        [x],
+        assert_fn=lambda x, y: np.testing.assert_allclose(x, y, rtol=RTOL, atol=ATOL),
+        py_mode="FAST_COMPILE",
+    )