DOC: fix tree_util.rst

saeta · Apr 8, 2021 · ffd1856 · ffd1856
1 parent 035c907
commit ffd1856
Showing 1 changed file with 26 additions and 2 deletions.
diff --git a/docs/jax.tree_util.rst b/docs/jax.tree_util.rst
@@ -1,7 +1,31 @@
 jax.tree_util package
 =====================
 
+.. currentmodule:: jax.tree_util
+
 .. automodule:: jax.tree_util
-    :members:
-    :undoc-members:
+
+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_multimap
+   tree_reduce
+   tree_structure
+   tree_transpose
+   tree_unflatten
+   treedef_children
+   treedef_is_leaf
+   treedef_tuple