Fix tests after applying JAX key-reuse checker. See:

- https://jax.readthedocs.io/en/latest/jax.experimental.key_reuse.html#experimental-key-reuse-checking
- https://jax.readthedocs.io/en/latest/_autosummary/jax.random.clone.html#jax.random.clone

PiperOrigin-RevId: 615600399

Merge pull request google#3682 from IvyZX:push

PiperOrigin-RevId: 605087622

Merge pull request google#3644 from chiamp:release

PiperOrigin-RevId: 600921456

Merge pull request google#3444 from levskaya:rc7.5

PiperOrigin-RevId: 577359040

Merge pull request google#3331 from levskaya:version_fix

PiperOrigin-RevId: 565109475

Merge pull request google#3326 from IvyZX:push

PiperOrigin-RevId: 564889489

Merge pull request google#3269 from chiamp:release

PiperOrigin-RevId: 555005843

Merge pull request google#3234 from chiamp:release

PiperOrigin-RevId: 552130592

Merge pull request google#3187 from google:0.7.0

PiperOrigin-RevId: 546394369

Merge pull request google#3155 from IvyZX:push

PiperOrigin-RevId: 542060958