Add SpecifyShape decorator to jax_funcify_Assert (#6062)

* Ignore SpecifyShape nodes in JAX mode

Co-authored-by: Martin Ingram <[email protected]>

Author: martiningram

pymc/sampling_jax.py

@@ -27,6 +27,7 @@
 from aesara.link.jax.dispatch import jax_funcify
 from aesara.raise_op import Assert
 from aesara.tensor import TensorVariable
+from aesara.tensor.shape import SpecifyShape
 from arviz.data.base import make_attrs
 
 from pymc import Model, modelcontext
@@ -38,6 +39,7 @@
 
 @jax_funcify.register(Assert)
 @jax_funcify.register(CheckParameterValue)
+@jax_funcify.register(SpecifyShape)
 def jax_funcify_Assert(op, **kwargs):
     # Jax does not allow assert whose values aren't known during JIT compilation
     # within it's JIT-ed code. Hence we need to make a simple pass through