Remove SpecifyShape from Assert JAX

Author: juanitorduz

pymc/sampling/jax.py

@@ -34,7 +34,6 @@
 from pytensor.raise_op import Assert
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.random.type import RandomType
-from pytensor.tensor.shape import SpecifyShape
 
 from pymc import Model, modelcontext
 from pymc.backends.arviz import find_constants, find_observations
@@ -62,7 +61,6 @@
 
 @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