2023-01-18 10:59:42 -08:00
|
|
|
``jax.tree_util`` module
|
|
|
|
|
========================
|
2019-05-14 21:00:27 -04:00
|
|
|
|
2021-04-08 15:59:55 -07:00
|
|
|
.. currentmodule:: jax.tree_util
|
|
|
|
|
|
2019-05-14 21:00:27 -04:00
|
|
|
.. automodule:: jax.tree_util
|
2021-04-08 15:59:55 -07:00
|
|
|
|
|
|
|
|
List of Functions
|
|
|
|
|
-----------------
|
|
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
|
|
|
|
Partial
|
|
|
|
|
all_leaves
|
|
|
|
|
build_tree
|
2024-05-16 18:34:54 -07:00
|
|
|
register_dataclass
|
2021-04-08 15:59:55 -07:00
|
|
|
register_pytree_node
|
|
|
|
|
register_pytree_node_class
|
2023-03-29 15:17:13 +00:00
|
|
|
register_pytree_with_keys
|
|
|
|
|
register_pytree_with_keys_class
|
2024-05-20 19:56:47 -07:00
|
|
|
register_static
|
2024-06-12 14:01:27 -07:00
|
|
|
tree_flatten_with_path
|
|
|
|
|
tree_leaves_with_path
|
|
|
|
|
tree_map_with_path
|
|
|
|
|
treedef_children
|
|
|
|
|
treedef_is_leaf
|
|
|
|
|
treedef_tuple
|
2024-07-31 06:40:53 -07:00
|
|
|
KeyEntry
|
|
|
|
|
KeyPath
|
2024-06-12 14:01:27 -07:00
|
|
|
keystr
|
|
|
|
|
|
|
|
|
|
Legacy APIs
|
|
|
|
|
-----------
|
|
|
|
|
These APIs are now accessed via :mod:`jax.tree`.
|
|
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
|
|
|
|
|
2021-04-08 15:59:55 -07:00
|
|
|
tree_all
|
|
|
|
|
tree_flatten
|
|
|
|
|
tree_leaves
|
|
|
|
|
tree_map
|
|
|
|
|
tree_reduce
|
|
|
|
|
tree_structure
|
|
|
|
|
tree_transpose
|
|
|
|
|
tree_unflatten
|
2019-05-14 21:00:27 -04:00
|
|
|
|
