Copybara import of the project:
--
d1935c71686533b2aa3d395d47395bf21457bffa by Neil Girdhar <mistersheik@gmail.com>:
Make FlatState a Mapping instead of a dict
Also, make traverse_util.flatten_dict accept mappings.
Fixes https://github.com/google/flax/issues/3879
PiperOrigin-RevId: 630158841
