Unbreak TFP on JAX OSS tests by pinning JAX version.
Unbreak TFP on JAX OSS tests by pinning JAX version.

c.f. https://github.com/google/jax/issues/10331

PiperOrigin-RevId: 442623798

tensorflow/probability,tensorflow/probability