style: miscellaneous style housekeeping (#92)

metaopt · Oct 5, 2023 · 7c4d71e · 7c4d71e
1 parent c1da9d0
commit 7c4d71e
Show file tree

Hide file tree

Showing 13 changed files with 158 additions and 102 deletions.
diff --git a/README.md b/README.md
@@ -156,8 +156,8 @@ optree.register_pytree_node(
     torch.Tensor,
     # (tensor) -> (children, metadata)
     flatten_func=lambda tensor: (
-        (tensor.cpu().numpy(),),
-        dict(dtype=tensor.dtype, device=tensor.device, requires_grad=tensor.requires_grad),
+        (tensor.cpu().detach().numpy(),),
+        {'dtype': tensor.dtype, 'device': tensor.device, 'requires_grad': tensor.requires_grad},
     ),
     # (metadata, children) -> tensor
     unflatten_func=lambda metadata, children: torch.tensor(children[0], **metadata),

diff --git a/docs/source/ops.rst b/docs/source/ops.rst
@@ -34,10 +34,10 @@ Tree Manipulation Functions
     tree_map_
     tree_map_with_path
     tree_map_with_path_
+    tree_replace_nones
     tree_transpose
     tree_broadcast_prefix
     broadcast_prefix
-    tree_replace_nones
     prefix_errors
 
 .. autofunction:: tree_flatten
@@ -51,10 +51,10 @@ Tree Manipulation Functions
 .. autofunction:: tree_map_
 .. autofunction:: tree_map_with_path
 .. autofunction:: tree_map_with_path_
+.. autofunction:: tree_replace_nones
 .. autofunction:: tree_transpose
 .. autofunction:: tree_broadcast_prefix
 .. autofunction:: broadcast_prefix
-.. autofunction:: tree_replace_nones
 .. autofunction:: prefix_errors
 
 ------
@@ -64,7 +64,6 @@ Tree Reduce Functions
 
 .. autosummary::
 
-    tree_replace_nones
     tree_reduce
     tree_sum
     tree_max

diff --git a/include/treespec.h b/include/treespec.h
@@ -249,8 +249,6 @@ class PyTreeSpec {
                                  const std::optional<py::function> &leaf_predicate,
                                  const std::string &registry_namespace);
 
-    [[nodiscard]] py::list FlattenUpToImpl(const py::handle &full_tree) const;
-
     template <bool NoneIsLeaf>
     static bool AllLeavesImpl(const py::iterable &iterable, const std::string &registry_namespace);
 

diff --git a/optree/__init__.py b/optree/__init__.py
@@ -96,10 +96,10 @@
     'tree_map_',
     'tree_map_with_path',
     'tree_map_with_path_',
+    'tree_replace_nones',
     'tree_transpose',
     'tree_broadcast_prefix',
     'broadcast_prefix',
-    'tree_replace_nones',
     'tree_reduce',
     'tree_sum',
     'tree_max',