Make FlatState a Mapping instead of a dict
* Add nnx.traversals.{flatten_mapping, unflatten_dict}
* These are modified from traverse_util.{flatten_dict,
unflatten_dict}
* They are annotated so that any future changes to the function
(e.g., changing it back to work with dicts only) triggers a type
error.
* Minor tweaks to imports:
* import from collections.abc instead of typing since the latter
imports are deprecated.
* Import from flax.typing instead of flax.core.scope since the
latter has been moved.
* Add from __future__ import annotations so that the annotations
work on Python 3.9.
* Minor tweaks to code:
* Annotate some private functions to make them easier to understand.
* When printing a type, print its __qualname__ since that's a bit
easier to read (str instead of <class 'str'>).
Fixes #3879
