Implement pad

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

@@ -0,0 +1,20 @@
+import jax.numpy as jnp
+
+from pytensor.link.jax.dispatch import jax_funcify
+from pytensor.tensor.pad import 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 = fixed_kwargs.get(pad_mode, {})
+    mode_kwargs = {kwarg: getattr(op, kwarg) for kwarg in expected_kwargs}
+
+    def pad(x, pad_width, *args):
+        print(args)
+        return jnp.pad(x, pad_width, mode=pad_mode, **mode_kwargs)
+
+    return pad

pytensor/tensor/__init__.py

@@ -130,6 +130,7 @@ def _get_vector_length_Constant(op: Op | Variable, var: Constant) -> int:
 from pytensor.tensor.extra_ops import *
 from pytensor.tensor.io import *
 from pytensor.tensor.math import *
+from pytensor.tensor.pad import pad
 from pytensor.tensor.shape import (
     reshape,
     shape,