Update

[ghstack-poisoned]

test/test_pytree.py

@@ -9,6 +9,7 @@ import sys
 import time
 import unittest
 from collections import defaultdict, deque, namedtuple, OrderedDict, UserDict
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from enum import auto
 from typing import Any, NamedTuple, Optional
@@ -22,6 +23,7 @@ from torch.testing._internal.common_utils import (
     parametrize,
     run_tests,
     subtest,
+    TEST_WITH_TORCHDYNAMO,
     TestCase,
 )
 
@@ -1472,8 +1474,23 @@ class TestCxxPytree(TestCase):
         if IS_FBCODE:
             raise unittest.SkipTest("C++ pytree tests are not supported in fbcode")
 
+    def assertEqualSpecs(
+        self,
+        spec1,
+        spec2,
+        msg: str | Callable[[str], str] | None = None,
+    ):
+        if TEST_WITH_TORCHDYNAMO:
+            # The Dynamo polyfill returns a pure Python class for PyTreeSpec.
+            # So we compare the type names and reprs instead because the types
+            # themselves won't be equal.
+            self.assertEqual(type(spec1).__name__, type(spec2).__name__, msg=msg)
+            self.assertEqual(repr(spec1), repr(spec2), msg=msg)
+        else:
+            self.assertEqual(spec1, spec2, msg=msg)
+
     def test_treespec_equality(self):
-        self.assertEqual(cxx_pytree.treespec_leaf(), cxx_pytree.treespec_leaf())
+        self.assertEqualSpecs(cxx_pytree.treespec_leaf(), cxx_pytree.treespec_leaf())
 
     def test_treespec_repr(self):
         # Check that it looks sane
@@ -1503,16 +1520,11 @@ class TestCxxPytree(TestCase):
         ],
     )
     def test_pytree_serialize(self, spec):
-        self.assertEqual(
-            spec,
-            cxx_pytree.tree_structure(
-                cxx_pytree.tree_unflatten([0] * spec.num_leaves, spec)
-            ),
-        )
-
         serialized_spec = cxx_pytree.treespec_dumps(spec)
         self.assertIsInstance(serialized_spec, str)
-        self.assertEqual(spec, cxx_pytree.treespec_loads(serialized_spec))
+
+        roundtrip_spec = cxx_pytree.treespec_loads(serialized_spec)
+        self.assertEqualSpecs(roundtrip_spec, spec)
 
     def test_pytree_serialize_namedtuple(self):
         python_pytree._register_namedtuple(
@@ -1538,7 +1550,7 @@ class TestCxxPytree(TestCase):
         spec = cxx_pytree.tree_structure(GlobalDummyType(0, 1))
         serialized_spec = cxx_pytree.treespec_dumps(spec)
         roundtrip_spec = cxx_pytree.treespec_loads(serialized_spec)
-        self.assertEqual(roundtrip_spec, spec)
+        self.assertEqualSpecs(roundtrip_spec, spec)
 
         class LocalDummyType:
             def __init__(self, x, y):
@@ -1554,7 +1566,7 @@ class TestCxxPytree(TestCase):
         spec = cxx_pytree.tree_structure(LocalDummyType(0, 1))
         serialized_spec = cxx_pytree.treespec_dumps(spec)
         roundtrip_spec = cxx_pytree.treespec_loads(serialized_spec)
-        self.assertEqual(roundtrip_spec, spec)
+        self.assertEqualSpecs(roundtrip_spec, spec)
 
 
 instantiate_parametrized_tests(TestGenericPytree)

torch/_dynamo/polyfills/pytree.py

@@ -346,8 +346,10 @@ if python_pytree._cxx_pytree_dynamo_traceable:
             assert callable(self._unflatten_func)
             return self._unflatten_func(self._metadata, subtrees)
 
-    def _is_pytreespec_instance(obj: Any, /) -> TypeIs[PyTreeSpec]:
-        return isinstance(obj, PyTreeSpec)
+    def _is_pytreespec_instance(
+        obj: Any, /
+    ) -> TypeIs[PyTreeSpec | python_pytree.TreeSpec]:
+        return isinstance(obj, (PyTreeSpec, python_pytree.TreeSpec))
 
     @substitute_in_graph(  # type: ignore[arg-type]
         optree.treespec_leaf,
@@ -550,7 +552,7 @@ if python_pytree._cxx_pytree_dynamo_traceable:
     def tree_unflatten(treespec: PyTreeSpec, leaves: Iterable[Any]) -> PyTree:
         if not _is_pytreespec_instance(treespec):
             raise TypeError(
-                f"tree_unflatten(leaves, treespec): Expected `treespec` to be instance of "
+                f"Expected `treespec` to be an instance of "
                 f"PyTreeSpec but got item of type {type(treespec)}."
             )
         return treespec.unflatten(leaves)

torch/utils/_cxx_pytree.py

@@ -13,6 +13,7 @@ collection support for PyTorch APIs.
 """
 
 import functools
+import sys
 import types
 from collections.abc import Callable, Iterable, Mapping
 from typing import Any, Optional, overload, TypeVar, Union
@@ -266,8 +267,21 @@ def _private_register_pytree_node(
         )
 
 
-def _is_pytreespec_instance(obj: Any, /) -> TypeIs[TreeSpec]:
-    return isinstance(obj, TreeSpec)
+def _is_pytreespec_instance(
+    obj: Any,
+    /,
+) -> TypeIs[Union[TreeSpec, python_pytree.PyTreeSpec]]:
+    if isinstance(obj, (TreeSpec, python_pytree.PyTreeSpec)):
+        return True
+    if "torch._dynamo.polyfills.pytree" in sys.modules:
+        # The PyTorch Dynamo pytree module is not always available, so we check if it is loaded.
+        # If the PyTorch Dynamo pytree module is loaded, we can check if the treespec
+        # is an instance of the PyTorch Dynamo TreeSpec class.
+        import torch._dynamo.polyfills.pytree as dynamo_pytree
+
+        if isinstance(obj, dynamo_pytree.PyTreeSpec):
+            return True
+    return False
 
 
 def treespec_leaf() -> TreeSpec:
@@ -394,7 +408,12 @@ def tree_unflatten(leaves: Iterable[Any], treespec: TreeSpec) -> PyTree:
         The reconstructed pytree, containing the ``leaves`` placed in the structure described by
         ``treespec``.
     """
-    return optree.tree_unflatten(treespec, leaves)  # type: ignore[arg-type]
+    if not _is_pytreespec_instance(treespec):
+        raise TypeError(
+            f"Expected `treespec` to be an instance of "
+            f"PyTreeSpec but got item of type {type(treespec)}."
+        )
+    return treespec.unflatten(leaves)
 
 
 def tree_iter(
@@ -973,7 +992,7 @@ def treespec_dumps(treespec: TreeSpec, protocol: Optional[int] = None) -> str:
     """Serialize a treespec to a JSON string."""
     if not _is_pytreespec_instance(treespec):
         raise TypeError(
-            f"treespec_dumps(treespec): Expected `treespec` to be instance of "
+            f"Expected `treespec` to be instance of "
             f"PyTreeSpec but got item of type {type(treespec)}."
         )
 

torch/utils/_pytree.py

@@ -20,6 +20,7 @@ import functools
 import importlib
 import importlib.metadata
 import json
+import sys
 import threading
 import types
 import warnings
@@ -36,14 +37,19 @@ from typing import (
     Optional,
     overload,
     Protocol,
+    TYPE_CHECKING,
     TypeVar,
     Union,
 )
-from typing_extensions import deprecated, NamedTuple, Self, TypeAlias
+from typing_extensions import deprecated, NamedTuple, Self, TypeAlias, TypeIs
 
 from torch.torch_version import TorchVersion as _TorchVersion
 
 
+if TYPE_CHECKING:
+    import torch.utils._cxx_pytree as cxx_pytree
+
+
 __all__ = [
     "PyTree",
     "Context",
@@ -249,9 +255,9 @@ def register_pytree_node(
         return
 
     if _cxx_pytree_imported:
-        from . import _cxx_pytree as cxx
+        import torch.utils._cxx_pytree as cxx_pytree
 
-        cxx._private_register_pytree_node(
+        cxx_pytree._private_register_pytree_node(
             cls,
             flatten_fn,
             unflatten_fn,
@@ -1363,6 +1369,45 @@ def treespec_dict(
     return TreeSpec(dict, list(dct.keys()), list(dct.values()))
 
 
+def _is_pytreespec_instance(
+    obj: Any,
+) -> TypeIs[Union[TreeSpec, "cxx_pytree.PyTreeSpec"]]:
+    if isinstance(obj, TreeSpec):
+        return True
+    if "torch.utils._cxx_pytree" in sys.modules:
+        # The C++ pytree module is not always available, so we check if it is loaded.
+        # If the C++ pytree module is loaded, we can check if the treespec
+        # is an instance of the C++ TreeSpec class.
+        import torch.utils._cxx_pytree as cxx_pytree
+
+        if isinstance(obj, cxx_pytree.PyTreeSpec):
+            return True
+    if "torch._dynamo.polyfills.pytree" in sys.modules:
+        # The PyTorch Dynamo pytree module is not always available, so we check if it is loaded.
+        # If the PyTorch Dynamo pytree module is loaded, we can check if the treespec
+        # is an instance of the PyTorch Dynamo TreeSpec class.
+        import torch._dynamo.polyfills.pytree as dynamo_pytree
+
+        if isinstance(obj, dynamo_pytree.PyTreeSpec):
+            return True
+    return False
+
+
+def _ensure_python_treespec_instance(
+    treespec: Union[TreeSpec, "cxx_pytree.PyTreeSpec"],
+) -> TreeSpec:
+    if isinstance(treespec, TreeSpec):
+        return treespec
+
+    if not _is_pytreespec_instance(treespec):
+        raise TypeError(
+            f"Expected `treespec` to be an instance of "
+            f"PyTreeSpec but got item of type {type(treespec)}."
+        )
+    dummy_tree = treespec.unflatten([0] * treespec.num_leaves)
+    return tree_structure(dummy_tree)
+
+
 def tree_flatten(
     tree: PyTree,
     is_leaf: Optional[Callable[[PyTree], bool]] = None,
@@ -1393,10 +1438,10 @@ def tree_unflatten(leaves: Iterable[Any], treespec: TreeSpec) -> PyTree:
     """Given a list of values and a TreeSpec, builds a pytree.
     This is the inverse operation of `tree_flatten`.
     """
-    if not isinstance(treespec, TreeSpec):
+    if not _is_pytreespec_instance(treespec):
         raise TypeError(
-            f"tree_unflatten(leaves, treespec): Expected `treespec` to be "
-            f"instance of TreeSpec but got item of type {type(treespec)}.",
+            f"Expected `treespec` to be an instance of "
+            f"PyTreeSpec but got item of type {type(treespec)}."
         )
     return treespec.unflatten(leaves)
 
@@ -1827,34 +1872,30 @@ def _broadcast_to_and_flatten(
     treespec: TreeSpec,
     is_leaf: Optional[Callable[[PyTree], bool]] = None,
 ) -> Optional[list[Any]]:
-    if not isinstance(treespec, TreeSpec):
-        raise AssertionError("treespec must be a TreeSpec")
+    def broadcast_prefix(
+        prefix_tree: PyTree,
+        full_tree: PyTree,
+        is_leaf: Optional[Callable[[PyTree], bool]] = None,
+    ) -> list[Any]:
+        result: list[Any] = []
 
-    if tree_is_leaf(tree, is_leaf=is_leaf):
-        return [tree] * treespec.num_leaves
-    if treespec.is_leaf():
+        def add_leaves(x: Any, subtree: PyTree) -> None:
+            subtreespec = tree_structure(subtree, is_leaf=is_leaf)
+            result.extend([x] * subtreespec.num_leaves)
+
+        tree_map_(
+            add_leaves,
+            prefix_tree,
+            full_tree,
+            is_leaf=is_leaf,
+        )
+        return result
+
+    full_tree = tree_unflatten([0] * treespec.num_leaves, treespec)
+    try:
+        return broadcast_prefix(tree, full_tree, is_leaf=is_leaf)
+    except ValueError:
         return None
-    node_type = _get_node_type(tree)
-    if node_type != treespec.type:
-        return None
-
-    flatten_fn = SUPPORTED_NODES[node_type].flatten_fn
-    child_pytrees, context = flatten_fn(tree)
-
-    # Check if the Node is different from the spec
-    if len(child_pytrees) != treespec.num_children or context != treespec._context:
-        return None
-
-    # Recursively flatten the children
-    result: list[Any] = []
-    for child, child_spec in zip(child_pytrees, treespec._children, strict=True):
-        flat = _broadcast_to_and_flatten(child, child_spec, is_leaf=is_leaf)
-        if flat is not None:
-            result += flat
-        else:
-            return None
-
-    return result
 
 
 @dataclasses.dataclass
@@ -1968,11 +2009,7 @@ _SUPPORTED_PROTOCOLS[1] = _ProtocolFn(_treespec_to_json, _json_to_treespec)
 
 
 def treespec_dumps(treespec: TreeSpec, protocol: Optional[int] = None) -> str:
-    if not isinstance(treespec, TreeSpec):
-        raise TypeError(
-            f"treespec_dumps(treespec, protocol): Expected `treespec` to be instance of "
-            f"TreeSpec but got item of type {type(treespec)}.",
-        )
+    treespec = _ensure_python_treespec_instance(treespec)
 
     if protocol is None:
         protocol = DEFAULT_TREESPEC_SERIALIZATION_PROTOCOL