mirror of
https://github.com/ROCm/jax.git
synced 2025-04-18 12:56:07 +00:00
31 lines
465 B
ReStructuredText
31 lines
465 B
ReStructuredText
jax.tree_util package
|
|
=====================
|
|
|
|
.. 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
|
|
|