BUGFix: Wrap roll as OpFromGraph with custom JAX dispatch to jnp.roll

pytensor/link/jax/dispatch/subtensor.py

@@ -1,4 +1,7 @@
+import jax.numpy as jnp
+
 from pytensor.link.jax.dispatch.basic import jax_funcify
+from pytensor.tensor.basic import Roll
 from pytensor.tensor.subtensor import (
     AdvancedIncSubtensor,
     AdvancedIncSubtensor1,
@@ -99,3 +102,13 @@ def makeslice(*x):
         return slice(*x)
 
     return makeslice
+
+
+@jax_funcify.register(Roll)
+def jax_funcify_Roll(op, **kwargs):
+    axis = op.axis
+
+    def roll(x, shift):
+        return jnp.roll(x, shift, axis=axis)
+
+    return roll

pytensor/tensor/basic.py

@@ -2865,9 +2865,21 @@ def roll(x, shift, axis=None):
     # List of slices describing the back half [:, :, :shift, :]
     end_slice = slice(0, -shift)
     end_list = [allslice] * axis + [end_slice] + [allslice] * (_x.ndim - axis - 1)
-    return join(
-        axis, _x.__getitem__(tuple(front_list)), _x.__getitem__(tuple(end_list))
-    )
+    out = join(axis, _x.__getitem__(tuple(front_list)), _x.__getitem__(tuple(end_list)))
+    return Roll(inputs=[_x, shift], outputs=[out], axis=axis)(_x, shift)
+
+
+class Roll(OpFromGraph):
+    """
+    Wrapper Op for roll graphs, enabling custom dispatch in JAX/Numba backends.
+    """
+
+    def __init__(self, *args, axis, **kwargs):
+        self.axis = axis
+        super().__init__(*args, **kwargs, strict=True)
+
+    def __str__(self):
+        return f"Roll{{axis={self.axis}}}"
 
 
 def stack(tensors: Sequence["TensorLike"], axis: int = 0):

tests/link/jax/test_tensor_basic.py

@@ -10,6 +10,7 @@
 from jax import errors
 
 import pytensor
+import pytensor.tensor as pt
 import pytensor.tensor.basic as ptb
 from pytensor.configdefaults import config
 from pytensor.tensor.type import iscalar, matrix, scalar, vector
@@ -233,3 +234,30 @@ def test_tri_nonconcrete():
         ),
     ):
         compare_jax_and_py([m, n, k], [out], [m_test_value, n_test_value, k_test_value])
+
+
+def test_jax_roll():
+    # Test roll with dynamic shift (the main bug fix)
+    x = pt.dmatrix("x")
+    shift = pt.iscalar("shift")
+    data = np.arange(16, dtype=float).reshape(4, 4)
+
+    # axis=0, dynamic shift
+    out = pt.roll(x, shift=shift, axis=0)
+    compare_jax_and_py([x, shift], [out], [data, 2])
+
+    # axis=1, dynamic shift
+    out = pt.roll(x, shift=shift, axis=1)
+    compare_jax_and_py([x, shift], [out], [data, 3])
+
+    # negative shift
+    out = pt.roll(x, shift=shift, axis=0)
+    compare_jax_and_py([x, shift], [out], [data, -1])
+
+    # axis=None (flatten then roll)
+    out = pt.roll(x, shift=shift, axis=None)
+    compare_jax_and_py([x, shift], [out], [data, 2])
+
+    # shift larger than axis size
+    out = pt.roll(x, shift=shift, axis=0)
+    compare_jax_and_py([x, shift], [out], [data, 10])