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_pytree_node
   register_pytree_node_class
   tree_all
   tree_flatten
   tree_leaves
   tree_map
   tree_reduce
   tree_structure
   tree_transpose
   tree_unflatten
   treedef_children
   treedef_is_leaf
   treedef_tuple