Making tests accommodate both FrozenDict and regular dict variables, as Flax will default to returning regular dicts in the future.
To toggle behavior between returning FrozenDict and regular dict variables at runtime, run flax.config.update('flax_return_frozendict', <BOOLEAN_VALUE>) in your code. Alternatively, the environment variable flax_return_frozendict (found here: https://github.com/google/flax/blob/main/flax/configurations.py#L121) can be directly modified in the Flax source code.
Making tests accommodate both FrozenDict and regular dict variables, as Flax will default to returning regular dicts in the future.
To toggle behavior between returning FrozenDict and regular dict variables at runtime, run
flax.config.update('flax_return_frozendict', <BOOLEAN_VALUE>)in your code. Alternatively, the environment variableflax_return_frozendict(found here: https://github.com/google/flax/blob/main/flax/configurations.py#L121) can be directly modified in the Flax source code.