graph
------------------------

.. automodule:: flax.nnx
.. currentmodule:: flax.nnx


.. autofunction:: split
.. autofunction:: merge
.. autofunction:: update
.. autofunction:: pop
.. autofunction:: state
.. autofunction:: variables
.. autofunction:: graph
.. autofunction:: graphdef
.. autofunction:: iter_graph
.. autofunction:: recursive_map
.. autofunction:: clone
.. autofunction:: call
.. autofunction:: set_metadata
.. autofunction:: cached_partial

.. autoclass:: GraphDef
  :members:

.. autoclass:: UpdateContext
  :members:

.. autofunction:: update_context
.. autofunction:: current_update_context

.. autofunction:: find_duplicates
.. autofunction:: pure
.. autofunction:: as_immutable_vars
.. autofunction:: as_mutable_vars
.. autofunction:: as_hijax_vars
.. autofunction:: as_pytree_vars
.. autofunction:: as_ref_vars
.. autofunction:: as_array_vars
.. autofunction:: flatten
.. autofunction:: unflatten
