Fix wrap_jax when there is a mix of statically known and unknown shapes

Fix wrap_jax when there is a mix of statically known and unknown shapes #3491