Fix wrap_jax when there is a mix of statically known and unknown shapes
#1722
+16
−3
Loading