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

Merged

1 commit merged intomainfrom test_614102332

Mar 14, 2024

+30-20

Commits on Mar 14, 2024

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

chiamp
authored and
Flax Authors
committed