jax.tree_util.rst

jax.tree_util module

.. currentmodule:: jax.tree_util

.. automodule:: jax.tree_util

List of Functions

.. autosummary::
   :toctree: _autosummary

   Partial
   all_leaves
   build_tree
   register_dataclass
   register_pytree_node
   register_pytree_node_class
   register_pytree_with_keys
   register_pytree_with_keys_class
   register_static
   tree_flatten_with_path
   tree_leaves_with_path
   tree_map_with_path
   treedef_children
   treedef_is_leaf
   treedef_tuple
   keystr

Legacy APIs

These APIs are now accessed via :mod:`jax.tree`.

.. autosummary::
   :toctree: _autosummary

   tree_all
   tree_flatten
   tree_leaves
   tree_map
   tree_reduce
   tree_structure
   tree_transpose
   tree_unflatten