Update

[ghstack-poisoned]
pytorch · XuehaiPan · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024
commit 5bbc7470f8bfc1e8b680c5e2f02ae73cf6745bd6
diff --git a/torch/_dynamo/polyfills/pytree.py b/torch/_dynamo/polyfills/pytree.py
@@ -191,6 +191,15 @@ def helper(node: PyTree, leaves: list[Any]) -> PyTreeSpec:
 
     __all__ += ["tree_flatten"]
 
+    @substitute_in_graph(cxx_pytree.tree_structure, can_constant_fold_through=False)  # type: ignore[arg-type]
+    def tree_structure(
+        tree: PyTree,
+        is_leaf: Callable[[PyTree], bool] | None = None,
+    ) -> PyTreeSpec:
+        return tree_flatten(tree, is_leaf=is_leaf)[1]  # type: ignore[return-value]
+
+    __all__ += ["tree_structure"]
+
     @substitute_in_graph(cxx_pytree.tree_unflatten, can_constant_fold_through=False)  # type: ignore[arg-type]
     def tree_unflatten(leaves: Iterable[Any], treespec: PyTreeSpec) -> PyTree:
         if not isinstance(treespec, PyTreeSpec):