flax.traverse_util package
============================

.. currentmodule:: flax.traverse_util

.. automodule:: flax.traverse_util

Traversal objects
-------------------

.. autoclass:: Traversal
    :members:

.. autoclass:: TraverseId
    :members:

.. autoclass:: TraverseMerge
    :members:

.. autoclass:: TraverseCompose
    :members:

.. autoclass:: TraverseFilter
    :members:

.. autoclass:: TraverseAttr
    :members:

.. autoclass:: TraverseItem
    :members:

.. autoclass:: TraverseEach
    :members:

.. autoclass:: TraverseTree
    :members:


Dict utils
------------

.. autofunction:: flatten_dict

.. autofunction:: unflatten_dict

.. autofunction:: path_aware_map


Model parameter traversal
--------------------------

.. autoclass:: ModelParamTraversal
    :members: __init__