Started failing in 0.4.36: https://github.com/jax-ml/jax/issues/26888

Skip failing JAX test

Started failing in 0.4.36 and seems to be fixed in 0.5.1