Announcement: Flax will move from returning FrozenDicts to regular dicts

[chiamp]

Hi everyone,

Flax will switch from returning FrozenDicts to returning regular dicts when calling the .init, .init_with_output and .apply Module methods. This has been a requested feature by users to simplify working with Flax params. For additional context, the original discussion can be found here.

Please refer to the migration guide to see how to migrate your code since this will be a breaking change.

The change is set to land next Monday, July 17th, 2023. As mentioned in the migration guide, a temporary feature flag flax_return_frozendict can be flipped to have Flax continue to return FrozenDicts at runtime. To do this, run flax.config.update('flax_return_frozendict', True) in your code. Please note, however, that this feature flag will be removed in the near future and is intended to aid in the migration of your code, and not intended as a permanent solution.

Update:
As of July 19th, 2023, Flax now defaults to returning regular dicts from version 0.7.1 onward. Also see #3193.