Update Oryx to be compatible with JAX's https://github.com/google/jax/pull/4008
Update Oryx to be compatible with JAX's https://github.com/google/jax/pull/4008

PiperOrigin-RevId: 337538591

tensorflow/probability,tensorflow/probability