flax

b94311e6 - Fix _is_valid_rngs for frozendict of rngs.