fix(jax): Specify shape should ignore None axes

Author: aseyboldt

pytensor/link/jax/dispatch/shape.py

@@ -96,12 +96,11 @@ def shape_i(x):
 def jax_funcify_SpecifyShape(op, node, **kwargs):
     def specifyshape(x, *shape):
         assert x.ndim == len(shape)
-        assert x.shape == tuple(shape), (
-            "got shape",
-            x.shape,
-            "expected",
-            shape,
-        )
+        for actual, expected in zip(x.shape, shape):
+            if expected is None:
+                continue
+            if actual != expected:
+                raise ValueError(f"Invalid shape: Expected {shape} but got {x.shape}")
         return x
 
     return specifyshape

tests/link/jax/test_shape.py

@@ -25,7 +25,7 @@ def test_jax_shape_ops():
 
 def test_jax_specify_shape():
     in_at = at.matrix("in")
-    x = at.specify_shape(in_at, (4, 5))
+    x = at.specify_shape(in_at, (4, None))
     x_fg = FunctionGraph([in_at], [x])
     compare_jax_and_py(x_fg, [np.ones((4, 5)).astype(config.floatX)])