rocm_jax/docs/jax.tree_util.rst
Tom Hennigan ed073aa6c9 Add jax.tree_util.tree_leaves_with_path(tree).
PiperOrigin-RevId: 539609052
2023-06-12 04:13:37 -07:00

37 lines
618 B
ReStructuredText

``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
register_pytree_with_keys
register_pytree_with_keys_class
tree_all
tree_flatten
tree_flatten_with_path
tree_leaves
tree_leaves_with_path
tree_map
tree_map_with_path
tree_reduce
tree_structure
tree_transpose
tree_unflatten
treedef_children
treedef_is_leaf
treedef_tuple
keystr