fix(ops): fix tests

metaopt · Feb 25, 2024 · 11dd785 · 11dd785
1 parent ab02dba
commit 11dd785
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/optree/ops.py b/optree/ops.py
@@ -362,6 +362,17 @@ def tree_flatten_with_typed_path(
         element is a list of leaf values and the last element is a treespec representing the
         structure of the pytree.
     """  # pylint: disable=line-too-long
+    leaves, treespec = _C.flatten(tree, is_leaf, none_is_leaf, namespace)
+    return treespec.typed_paths(), leaves, treespec
+
+
+def tree_flatten_with_accessor(
+    tree: PyTree[T],
+    is_leaf: Callable[[T], bool] | None = None,
+    *,
+    none_is_leaf: bool = False,
+    namespace: str = '',
+) -> tuple[list[PyTreeAccessor], list[T], PyTreeSpec]:
     leaves, treespec = _C.flatten(tree, is_leaf, none_is_leaf, namespace)
     return (
         [