google / aqt

Apache License 2.0
266 stars 27 forks source link

Fix cases of prng key reuse in aqt #515

Closed copybara-service[bot] closed 8 months ago

copybara-service[bot] commented 8 months ago

Fix cases of prng key reuse in aqt

Detected via running tests with configuration jax_enable_key_reuse_checks=true, available since https://github.com/google/jax/pull/19795