diff --git a/python/celerite2/pymc/__pycache__/ops.cpython-312.pyc b/python/celerite2/pymc/__pycache__/ops.cpython-312.pyc
new file mode 100644
index 0000000..bb5de2c
Binary files /dev/null and b/python/celerite2/pymc/__pycache__/ops.cpython-312.pyc differ
diff --git a/python/celerite2/pymc/ops.py b/python/celerite2/pymc/ops.py
index cce6928..da45f64 100644
--- a/python/celerite2/pymc/ops.py
+++ b/python/celerite2/pymc/ops.py
@@ -19,11 +19,18 @@
 import pytensor
 import pytensor.tensor as pt
 from pytensor.graph import basic, op
-from pytensor.link.jax.dispatch import jax_funcify
 
 import celerite2.backprop as backprop
 import celerite2.driver as driver
 
+try:
+    from pytensor.link.jax.dispatch import jax_funcify
+except ImportError:  # pragma: no cover - jax is an optional dependency
+    # pytensor.link.jax.dispatch imports jax, which is not a required
+    # dependency of the PyMC backend. The conversion defined below is then
+    # simply not registered.
+    jax_funcify = None
+
 
 def _resize_or_set(outputs, n, shape):
     if outputs[n][0] is None:
@@ -160,7 +167,6 @@ def grad(self, inputs, gradients):
 
 
 # JAX conversion for PyTensor JAX linker -------------------------------------
-@jax_funcify.register(_CeleriteOp)
 def _jax_funcify_celerite(op, node, **kwargs):
     """Map celerite2 PyTensor ops to their JAX counterparts."""
 
@@ -224,3 +230,7 @@ def general_matmul_upper_fwd(t1, t2, c, U, V, Y):
         raise NotImplementedError(
             f"No JAX conversion registered for {op.name}"
         )
+
+
+if jax_funcify is not None:
+    jax_funcify.register(_CeleriteOp)(_jax_funcify_celerite)