Update

[ghstack-poisoned]

CODEOWNERS

@@ -195,6 +195,7 @@ torch/backends/cudnn/ @eqy @syed-ahmed @Aidyn-A
 /torch/utils/_pytree.py @XuehaiPan
 /torch/utils/_cxx_pytree.py @XuehaiPan
 /torch/utils/pytree/ @XuehaiPan
+/torch/pytree.py @XuehaiPan
 /torch/_dynamo/polyfills/pytree.py @XuehaiPan
 
 # Relating to libtorch ABI

docs/source/pytorch-api.md

@@ -60,6 +60,7 @@ torch.special <special>
 torch.overrides
 torch.nativert <nativert>
 torch.package <package>
+torch.pytree <pytree>
 profiler
 nn.init
 nn.attention
@@ -77,6 +78,7 @@ sparse
 storage
 torch.testing <testing>
 torch.utils <utils>
+torch.utils.pytree
 torch.utils.benchmark <benchmark_utils>
 torch.utils.checkpoint <checkpoint>
 torch.utils.cpp_extension <cpp_extension>

docs/source/pytree.rst

@@ -0,0 +1,7 @@
+torch.pytree
+============
+
+.. currentmodule:: torch.pytree
+
+.. automodule:: torch.pytree
+    :members:

docs/source/torch.utils.pytree.rst

@@ -0,0 +1,7 @@
+torch.utils.pytree
+==================
+
+.. currentmodule:: torch.utils.pytree
+
+.. automodule:: torch.utils.pytree
+    :members:

mypy-strict.ini

@@ -29,6 +29,7 @@ files =
     benchmarks/instruction_counts,
     tools,
     torch/profiler/_memory_profiler.py,
+    torch/utils/pytree/__init__.py,
     torch/utils/_pytree.py,
     torch/utils/_cxx_pytree.py,
     torch/utils/benchmark/utils/common.py,

test/allowlist_for_publicAPI.json

@@ -687,6 +687,28 @@
     "kineto_available",
     "record_function"
   ],
+  "torch.pytree": [
+    "PyTreeSpec",
+    "register_node",
+    "all",
+    "all_only",
+    "any",
+    "any_only",
+    "flatten",
+    "iter",
+    "leaves",
+    "map",
+    "map_",
+    "map_only",
+    "map_only_",
+    "structure",
+    "is_namedtuple",
+    "is_namedtuple_class",
+    "is_namedtuple_instance",
+    "is_structseq",
+    "is_structseq_class",
+    "is_structseq_instance"
+  ],
   "torch.quantization": [
     "ABC",
     "DeQuantStub",

test/dynamo/test_flat_apply.py

@@ -147,8 +147,8 @@ class GraphModule(torch.nn.Module):
 
         t: "f32[10]" = l_x_ + l_y_
 
-        trace_point_tensor_spec : torch.utils._pytree.TreeSpec = self.trace_point_tensor_spec
-        trace_point_tensor_input_spec : torch.utils._pytree.TreeSpec = self.trace_point_tensor_input_spec
+        trace_point_tensor_spec : torch.utils.pytree.python.PyTreeSpec = self.trace_point_tensor_spec
+        trace_point_tensor_input_spec : torch.utils.pytree.python.PyTreeSpec = self.trace_point_tensor_input_spec
         res: "f32[10]" = torch.ops.higher_order.flat_apply(trace_point_tensor_spec, trace_point_tensor_input_spec, l_x_, l_y_, t);  trace_point_tensor_spec = trace_point_tensor_input_spec = l_x_ = l_y_ = t = None
         return (res,)
 """,  # NOQA: B950

test/dynamo/test_misc.py

@@ -40,6 +40,7 @@ import torch._inductor.test_case
 import torch.onnx.operators
 import torch.utils._pytree as python_pytree
 import torch.utils.cpp_extension
+import torch.utils.pytree as generic_pytree
 from torch import Tensor
 from torch._C import FileCheck
 from torch._dynamo import allow_in_graph
@@ -104,6 +105,7 @@ from torch.testing._internal.jit_utils import JitTestCase
 
 
 pytree_modules = {
+    "generic": generic_pytree,
     "python": python_pytree,
 }
 if python_pytree._cxx_pytree_dynamo_traceable:

test/export/test_export.py

@@ -8296,7 +8296,7 @@ graph():
     %to : [num_users=1] = call_function[target=operator.getitem](args = (%tree_flatten_spec, 0), kwargs = {})
     %sum_1 : [num_users=1] = call_function[target=torch.ops.aten.sum.default](args = (%to,), kwargs = {})
     %_spec_1 : [num_users=1] = get_attr[target=_spec_1]
-    %tree_unflatten : [num_users=1] = call_function[target=torch.utils._pytree.tree_unflatten](args = ((%sum_1,), %_spec_1), kwargs = {})
+    %tree_unflatten : [num_users=1] = call_function[target=torch.utils.pytree.python.tree_unflatten](args = ((%sum_1,), %_spec_1), kwargs = {})
     return tree_unflatten""",
         )
 

test/export/test_swap.py

@@ -249,11 +249,11 @@ def forward(self, x, y):
     _spec_0 = self._spec_0
     _spec_1 = self._spec_1
     _spec_4 = self._spec_4
-    tree_flatten = torch.utils._pytree.tree_flatten((x_1, y_1));  x_1 = y_1 = None
+    tree_flatten = torch.utils.pytree.python.tree_flatten((x_1, y_1));  x_1 = y_1 = None
     getitem = tree_flatten[0];  tree_flatten = None
     x = getitem[0]
     y = getitem[1];  getitem = None
-    tree_unflatten_1 = torch.utils._pytree.tree_unflatten([x, y], _spec_1);  x = y = _spec_1 = None
+    tree_unflatten_1 = torch.utils.pytree.python.tree_unflatten([x, y], _spec_1);  x = y = _spec_1 = None
     getitem_1 = tree_unflatten_1[0];  tree_unflatten_1 = None
     getitem_2 = getitem_1[0]
     getitem_3 = getitem_1[1];  getitem_1 = None
@@ -261,7 +261,7 @@ def forward(self, x, y):
     bar = self.bar(foo);  foo = None
     tree_flatten_spec_1 = torch.fx._pytree.tree_flatten_spec(bar, _spec_4);  bar = _spec_4 = None
     getitem_10 = tree_flatten_spec_1[0];  tree_flatten_spec_1 = None
-    tree_unflatten = torch.utils._pytree.tree_unflatten((getitem_10,), _spec_0);  getitem_10 = _spec_0 = None
+    tree_unflatten = torch.utils.pytree.python.tree_unflatten((getitem_10,), _spec_0);  getitem_10 = _spec_0 = None
     return tree_unflatten""",
         )
 
@@ -324,7 +324,7 @@ def forward(self, x, y):
     x, y, = fx_pytree.tree_flatten_spec(([x, y], {}), self._in_spec)
     _spec_0 = self._spec_0
     _spec_3 = self._spec_3
-    tree_unflatten = torch.utils._pytree.tree_unflatten([x, y], _spec_0);  x = y = _spec_0 = None
+    tree_unflatten = torch.utils.pytree.python.tree_unflatten([x, y], _spec_0);  x = y = _spec_0 = None
     getitem = tree_unflatten[0];  tree_unflatten = None
     getitem_1 = getitem[0]
     getitem_2 = getitem[1];  getitem = None

test/test_pytree.py

@@ -22,6 +22,7 @@ from torch.testing._internal.common_utils import (
     parametrize,
     run_tests,
     subtest,
+    TEST_WITH_TORCHDYNAMO,
     TestCase,
 )
 
@@ -52,6 +53,14 @@ class GlobalDummyType:
         self.x = x
         self.y = y
 
+    def __eq__(self, other):
+        if not isinstance(other, GlobalDummyType):
+            return NotImplemented
+        return self.x == other.x and self.y == other.y
+
+    def __hash__(self):
+        return hash((self.x, self.y))
+
 
 cxx_pytree.register_pytree_node(
     GlobalDummyType,
@@ -156,6 +165,44 @@ class TestGenericPytree(TestCase):
                             ),
                         )
 
+    @parametrize(
+        "modulename",
+        [
+            subtest("python", name="py"),
+            *([subtest("cxx", name="cxx")] if not IS_FBCODE else []),
+        ],
+    )
+    def test_public_api_import(self, modulename):
+        for use_cxx_pytree in [None, "", "0", *(["1"] if not IS_FBCODE else [])]:
+            env = os.environ.copy()
+            if use_cxx_pytree is not None:
+                env["PYTORCH_USE_CXX_PYTREE"] = str(use_cxx_pytree)
+            else:
+                env.pop("PYTORCH_USE_CXX_PYTREE", None)
+            for statement in (
+                f"import torch.utils.pytree.{modulename}",
+                f"from torch.utils.pytree import {modulename}",
+                f"from torch.utils.pytree.{modulename} import tree_map",
+                f"import torch.utils.pytree; torch.utils.pytree.{modulename}",
+                f"import torch.utils.pytree; torch.utils.pytree.{modulename}.tree_map",
+            ):
+                try:
+                    subprocess.check_output(
+                        [sys.executable, "-c", statement],
+                        stderr=subprocess.STDOUT,
+                        # On Windows, opening the subprocess with the default CWD makes `import torch`
+                        # fail, so just set CWD to this script's directory
+                        cwd=os.path.dirname(os.path.realpath(__file__)),
+                        env=env,
+                    )
+                except subprocess.CalledProcessError as e:
+                    self.fail(
+                        msg=(
+                            f"Subprocess exception while attempting to run statement `{statement}`: "
+                            + e.output.decode("utf-8")
+                        )
+                    )
+
     @parametrize_pytree_module
     def test_register_pytree_node(self, pytree):
         class MyDict(UserDict):
@@ -838,12 +885,13 @@ class TestPythonPytree(TestCase):
         script = """
 import sys
 import torch
-import torch.utils._pytree
+assert "torch.utils.pytree" in sys.modules
 assert "torch.utils._pytree" in sys.modules
-if "torch.utils._cxx_pytree" in sys.modules:
-    raise RuntimeError("importing torch.utils._pytree should not import torch.utils._cxx_pytree")
-if "optree" in sys.modules:
-    raise RuntimeError("importing torch.utils._pytree should not import optree")
+if not torch.utils.pytree.PYTORCH_USE_CXX_PYTREE:
+    if "torch.utils._cxx_pytree" in sys.modules:
+        raise RuntimeError("importing torch.utils._pytree should not import torch.utils._cxx_pytree")
+    if "optree" in sys.modules:
+        raise RuntimeError("importing torch.utils._pytree should not import optree")
 """
         try:
             subprocess.check_output(
@@ -1490,6 +1538,25 @@ class TestCxxPytree(TestCase):
         if IS_FBCODE:
             raise unittest.SkipTest("C++ pytree tests are not supported in fbcode")
 
+    def assertEqual(self, x, y, *args, **kwargs):
+        x_typename, y_typename = type(x).__name__, type(y).__name__
+        if not ("treespec" in x_typename.lower() or "treespec" in y_typename.lower()):
+            super().assertEqual(x, y, *args, **kwargs)
+
+        # The Dynamo polyfill returns a polyfilled Python class for C++ PyTreeSpec instead of the
+        # C++ class. So we compare the type names and reprs instead because the types themselves
+        # won't be equal.
+        super().assertEqual(x_typename, y_typename, *args, **kwargs)
+        if not TEST_WITH_TORCHDYNAMO or type(x) is type(y):
+            super().assertEqual(x, y, *args, **kwargs)
+        else:
+            super().assertEqual(
+                x.unflatten(range(x.num_leaves)),
+                y.unflatten(range(y.num_leaves)),
+                *args,
+                **kwargs,
+            )
+
     def test_treespec_equality(self):
         self.assertEqual(cxx_pytree.treespec_leaf(), cxx_pytree.treespec_leaf())
 
@@ -1530,7 +1597,9 @@ class TestCxxPytree(TestCase):
 
         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.assertEqual(roundtrip_spec, spec)
 
     def test_pytree_serialize_namedtuple(self):
         python_pytree._register_namedtuple(
@@ -1563,6 +1632,14 @@ class TestCxxPytree(TestCase):
                 self.x = x
                 self.y = y
 
+            def __eq__(self, other):
+                if not isinstance(other, LocalDummyType):
+                    return NotImplemented
+                return self.x == other.x and self.y == other.y
+
+            def __hash__(self):
+                return hash((self.x, self.y))
+
         cxx_pytree.register_pytree_node(
             LocalDummyType,
             lambda dummy: ([dummy.x, dummy.y], None),

torch/__init__.py

@@ -2794,6 +2794,7 @@ if TYPE_CHECKING:
         _inductor as _inductor,
         _subclasses as _subclasses,
         onnx as onnx,
+        pytree as pytree,
     )
 
 else:
@@ -2803,6 +2804,7 @@ else:
         "_export",
         # ONNX must be imported after _dynamo, _ops, _subclasses, fx, func and jit
         "onnx",
+        "pytree",
     }
 
     def __getattr__(name):

torch/_dynamo/polyfills/pytree.py

@@ -1,3 +1,5 @@
+# Owner(s): ["module: pytree"]
+
 """
 Python polyfills for torch.utils.pytree
 """
@@ -23,6 +25,7 @@ from optree import (
 )
 
 import torch.utils._cxx_pytree as cxx_pytree  # noqa: F401
+import torch.utils._pytree as python_pytree
 from torch.utils._pytree import BUILTIN_TYPES, STANDARD_DICT_TYPES
 
 from ..decorators import substitute_in_graph
@@ -427,8 +430,8 @@ class PyTreeSpec:
         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]
@@ -698,7 +701,7 @@ def tree_structure(
 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/_dynamo/trace_rules.py

@@ -3454,6 +3454,7 @@ MOD_INLINELIST = [
     "torch.utils._python_dispatch",
     "torch.utils._pytree",
     "torch.utils.hooks",
+    "torch.utils.pytree",
 ]
 assert sorted(set(MOD_INLINELIST)) == MOD_INLINELIST
 MOD_INLINELIST = set(MOD_INLINELIST)

torch/_export/utils.py

@@ -39,13 +39,13 @@ from torch.utils._pytree import (
     _register_pytree_node,
     Context,
     FlattenFunc,
-    FromDumpableContextFn,
+    FromDumpableContextFunc,
     GetAttrKey,
     KeyPath,
     keystr,
     MappingKey,
     SequenceKey,
-    ToDumpableContextFn,
+    ToDumpableContextFunc,
     tree_flatten_with_path,
     UnflattenFunc,
 )
@@ -485,8 +485,8 @@ def register_dataclass_as_pytree_node(
     unflatten_fn: Optional[UnflattenFunc] = None,
     *,
     serialized_type_name: Optional[str] = None,
-    to_dumpable_context: Optional[ToDumpableContextFn] = None,
-    from_dumpable_context: Optional[FromDumpableContextFn] = None,
+    to_dumpable_context: Optional[ToDumpableContextFunc] = None,
+    from_dumpable_context: Optional[FromDumpableContextFunc] = None,
     return_none_fields: bool = False,
 ) -> None:
     assert dataclasses.is_dataclass(cls), (

torch/export/_swap.py

@@ -42,7 +42,7 @@ def _try_remove_connecting_pytrees(curr_module_node: torch.fx.Node) -> None:
         %foo : [num_users=1] = call_module[target=foo](args = (%getitem_1, %getitem_2), kwargs = {})
         %tree_flatten_spec : [num_users=1] = call_function[target=torch.fx._pytree.tree_flatten_spec](args = (%foo, %_spec_1), kwargs = {})
         %getitem_4 : [num_users=1] = call_function[target=operator.getitem](args = (%tree_flatten_spec, 0), kwargs = {})
-        %tree_unflatten_1 : [num_users=2] = call_function[target=torch.utils._pytree.tree_unflatten](args = ([%getitem_4], %_spec_2), kwargs = {})
+        %tree_unflatten_1 : [num_users=2] = call_function[target=torch.utils.pytree.python.tree_unflatten](args = ([%getitem_4], %_spec_2), kwargs = {})
         %getitem_5 : [num_users=1] = call_function[target=operator.getitem](args = (%tree_unflatten_1, 0), kwargs = {})
         %getitem_7 : [num_users=0] = call_function[target=operator.getitem](args = (%tree_unflatten_1, 1), kwargs = {})
         %getitem_6 : [num_users=1] = call_function[target=operator.getitem](args = (%getitem_5, 0), kwargs = {})
@@ -293,7 +293,7 @@ def _swap_module_helper(
             %y : [num_users=1] = placeholder[target=y]
 
             %_spec_0 : [num_users=1] = get_attr[target=_spec_0]
-            %tree_unflatten : [num_users=2] = call_function[target=torch.utils._pytree.tree_unflatten](args = ([%x, %y], %_spec_0), kwargs = {})
+            %tree_unflatten : [num_users=2] = call_function[target=torch.utils.pytree.python.tree_unflatten](args = ([%x, %y], %_spec_0), kwargs = {})
             %getitem : [num_users=2] = call_function[target=operator.getitem](args = (%tree_unflatten, 0), kwargs = {})
             %getitem_1 : [num_users=1] = call_function[target=operator.getitem](args = (%getitem, 0), kwargs = {})
             %getitem_2 : [num_users=1] = call_function[target=operator.getitem](args = (%getitem, 1), kwargs = {})

torch/fx/graph.py

@@ -88,7 +88,7 @@ _register_custom_builtin("NoneType", "NoneType = type(None)", type(None))
 _register_custom_builtin("torch", "import torch", torch)
 _register_custom_builtin("device", "from torch import device", torch.device)
 _register_custom_builtin("fx_pytree", "import torch.fx._pytree as fx_pytree", fx_pytree)
-_register_custom_builtin("pytree", "import torch.utils._pytree as pytree", pytree)
+_register_custom_builtin("pytree", "import torch.utils.pytree.python as pytree", pytree)
 
 
 def _is_magic(x: str) -> bool:

torch/pytree.py

@@ -0,0 +1,109 @@
+# Owner(s): ["module: pytree"]
+
+"""
+Contains utility functions for working with nested python data structures.
+
+A *pytree* is Python nested data structure. It is a tree in the sense that
+nodes are Python collections (e.g., list, tuple, dict) and the leaves are
+Python values. Furthermore, a pytree should not contain reference cycles.
+
+pytrees are useful for working with nested collections of Tensors. For example,
+one can use `map` to map a function over all Tensors inside some nested
+collection of Tensors and `leaves` to get a flat list of all Tensors
+inside some nested collection. pytrees are helpful for implementing nested
+collection support for PyTorch APIs.
+"""
+
+from __future__ import annotations
+
+from typing import Any as _Any, TYPE_CHECKING as _TYPE_CHECKING
+
+from torch.utils.pytree import (
+    is_namedtuple,
+    is_namedtuple_class,
+    is_namedtuple_instance,
+    is_structseq,
+    is_structseq_class,
+    is_structseq_instance,
+    PyTree,
+    PyTreeSpec,
+    register_pytree_node as register_node,
+    tree_all as all,
+    tree_all_only as all_only,
+    tree_any as any,
+    tree_any_only as any_only,
+    tree_flatten as flatten,
+    tree_iter as iter,
+    tree_leaves as leaves,
+    tree_map as map,
+    tree_map_ as map_,
+    tree_map_only as map_only,
+    tree_map_only_ as map_only_,
+    tree_structure as structure,
+    tree_unflatten as _tree_unflatten,
+)
+
+
+if _TYPE_CHECKING:
+    from collections.abc import Iterable
+
+
+__all__ = [
+    "PyTreeSpec",
+    "register_node",
+    "flatten",
+    "unflatten",
+    "iter",
+    "leaves",
+    "structure",
+    "map",
+    "map_",
+    "map_only",
+    "map_only_",
+    "all",
+    "any",
+    "all_only",
+    "any_only",
+    "is_namedtuple",
+    "is_namedtuple_class",
+    "is_namedtuple_instance",
+    "is_structseq",
+    "is_structseq_class",
+    "is_structseq_instance",
+]
+
+
+def unflatten(treespec: PyTreeSpec, leaves: Iterable[_Any]) -> PyTree:
+    """Reconstruct a pytree from the treespec and the leaves.
+
+    The inverse of :func:`flatten`.
+
+    >>> tree = {"b": (2, [3, 4]), "a": 1, "c": None, "d": 5}
+    >>> leaves, treespec = torch.pytree.flatten(tree)
+    >>> tree == torch.pytree.unflatten(treespec, leaves)
+    True
+
+    .. warning::
+
+        This function has a different signature than :func:`torch.utils.pytree.tree_unflatten`.
+        The ``treespec`` argument comes first to have a better :class:`functools.partial` support:
+
+        .. code-block:: python
+
+            import functools
+
+            unflatten_fn = functools.partial(torch.pytree.unflatten, treespec)
+            tree1 = unflatten_fn(leaves1)
+            tree2 = unflatten_fn(leaves2)
+
+    Args:
+        treespec (PyTreeSpec): The treespec to reconstruct.
+        leaves (iterable): The list of leaves to use for reconstruction. The list must match the
+            number of leaves of the treespec.
+
+    Returns:
+        The reconstructed pytree, containing the ``leaves`` placed in the structure described by
+        ``treespec``.
+    """
+    # pyrefly: ignore [bad-argument-type]
+    return _tree_unflatten(leaves, treespec)

torch/utils/__init__.py

@@ -11,6 +11,7 @@ from torch.utils import (
     data as data,
     deterministic as deterministic,
     hooks as hooks,
+    pytree as pytree,
 )
 from torch.utils.backend_registration import (
     generate_methods_for_privateuse1_backend,

torch/utils/_cxx_pytree.py

@@ -1,3 +1,5 @@
+# Owner(s): ["module: pytree"]
+
 """
 Contains utility functions for working with nested python data structures.
 
@@ -13,6 +15,7 @@ collection support for PyTorch APIs.
 """
 
 import functools
+import sys
 import types
 from collections.abc import Callable, Iterable, Mapping
 from typing import Any, overload, TypeAlias, TypeVar, Union
@@ -21,13 +24,21 @@ from typing_extensions import deprecated, Self, TypeIs
 import torch.utils._pytree as python_pytree
 from torch.torch_version import TorchVersion as _TorchVersion
 from torch.utils._pytree import (
+    Context,
+    DumpableContext,
+    FlattenFunc,
+    FlattenWithKeysFunc,
+    FromDumpableContextFunc,
     is_namedtuple,
     is_namedtuple_class,
     is_namedtuple_instance,
     is_structseq,
     is_structseq_class,
     is_structseq_instance,
-    KeyEntry,
+    KeyPath,
+    PyTree,
+    ToDumpableContextFunc,
+    UnflattenFunc,
 )
 
 
@@ -51,8 +62,8 @@ __all__ = [
     "FlattenFunc",
     "UnflattenFunc",
     "DumpableContext",
-    "ToDumpableContextFn",
-    "FromDumpableContextFn",
+    "ToDumpableContextFunc",
+    "FromDumpableContextFunc",
     "PyTreeSpec",
     "TreeSpec",
     "LeafSpec",
@@ -88,6 +99,9 @@ __all__ = [
 ]
 
 
+__name__ = "torch.utils.pytree.cxx"  # sets the __module__ attribute of all functions in this module
+
+
 # In-tree installation may have VCS-based versioning. Update the previous static version.
 python_pytree._optree_version = _TorchVersion(optree.__version__)  # type: ignore[attr-defined]
 
@@ -100,19 +114,8 @@ S = TypeVar("S")
 U = TypeVar("U")
 R = TypeVar("R")
 
-
 TreeSpec: TypeAlias = PyTreeSpec
-
-Context = Any
-PyTree = Any
-FlattenFunc = Callable[[PyTree], tuple[list[Any], Context]]
-UnflattenFunc = Callable[[Iterable[Any], Context], PyTree]
-OpTreeUnflattenFunc = Callable[[Context, Iterable[Any]], PyTree]
-DumpableContext = Any  # Any json dumpable text
-ToDumpableContextFn = Callable[[Context], DumpableContext]
-FromDumpableContextFn = Callable[[DumpableContext], Context]
-KeyPath = tuple[KeyEntry, ...]
-FlattenWithKeysFunc = Callable[[PyTree], tuple[list[tuple[KeyEntry, Any]], Any]]
+OpTreeUnflattenFunc: TypeAlias = Callable[[Context, Iterable[Any]], PyTree]
 
 
 def _reverse_args(func: UnflattenFunc) -> OpTreeUnflattenFunc:
@@ -129,8 +132,8 @@ def register_pytree_node(
     unflatten_fn: UnflattenFunc,
     *,
     serialized_type_name: str | None = None,
-    to_dumpable_context: ToDumpableContextFn | None = None,
-    from_dumpable_context: FromDumpableContextFn | None = None,
+    to_dumpable_context: ToDumpableContextFunc | None = None,
+    from_dumpable_context: FromDumpableContextFunc | None = None,
     flatten_with_keys_fn: FlattenWithKeysFunc | None = None,
 ) -> None:
     """Register a container-like type as pytree node.
@@ -197,8 +200,8 @@ def _register_pytree_node(
     unflatten_fn: UnflattenFunc,
     *,
     serialized_type_name: str | None = None,
-    to_dumpable_context: ToDumpableContextFn | None = None,
-    from_dumpable_context: FromDumpableContextFn | None = None,
+    to_dumpable_context: ToDumpableContextFunc | None = None,
+    from_dumpable_context: FromDumpableContextFunc | None = None,
 ) -> None:
     """Register a container-like type as pytree node for the C++ pytree only.
 
@@ -248,8 +251,8 @@ def _private_register_pytree_node(
     unflatten_fn: UnflattenFunc,
     *,
     serialized_type_name: str | None = None,
-    to_dumpable_context: ToDumpableContextFn | None = None,
-    from_dumpable_context: FromDumpableContextFn | None = None,
+    to_dumpable_context: ToDumpableContextFunc | None = None,
+    from_dumpable_context: FromDumpableContextFunc | None = None,
 ) -> None:
     """This is an internal function that is used to register a pytree node type
     for the C++ pytree only. End-users should use :func:`register_pytree_node`
@@ -266,8 +269,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 +410,15 @@ 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):
+        if not _is_pytreespec_instance(leaves):
+            raise TypeError(
+                f"Expected `treespec` to be an instance of "
+                f"PyTreeSpec but got item of type {type(treespec)}."
+            )
+        # Allow passing the PyTreeSpec instance as the first argument
+        leaves, treespec = treespec, leaves
+    return treespec.unflatten(leaves)
 
 
 def tree_iter(
@@ -959,8 +983,9 @@ def _broadcast_to_and_flatten(
     is_leaf: Callable[[PyTree], bool] | None = None,
 ) -> list[Any] | None:
     if not _is_pytreespec_instance(treespec):
-        raise AssertionError(
-            f"_broadcast_to_and_flatten: Expected `treespec` to be instance of PyTreeSpec but got {type(treespec)}"
+        raise TypeError(
+            f"Expected `treespec` to be an instance of "
+            f"PyTreeSpec but got item of type {type(treespec)}."
         )
     full_tree = tree_unflatten([0] * treespec.num_leaves, treespec)
     try:
@@ -973,7 +998,7 @@ def treespec_dumps(treespec: TreeSpec, protocol: int | None = 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 an instance of "
             f"PyTreeSpec but got item of type {type(treespec)}."
         )
 
@@ -994,16 +1019,22 @@ def treespec_loads(serialized: str) -> TreeSpec:
     return treespec
 
 
-class _DummyLeaf:
+class _Asterisk(str):
+    __slots__ = ()
+
+    def __new__(cls) -> Self:
+        return super().__new__(cls, "*")
+
     def __repr__(self) -> str:
-        return "*"
+        return "*"  # no quotes
+
+
+_asterisk = _Asterisk()
+del _Asterisk
 
 
 def treespec_pprint(treespec: TreeSpec) -> str:
-    dummy_tree = tree_unflatten(
-        [_DummyLeaf() for _ in range(treespec.num_leaves)],
-        treespec,
-    )
+    dummy_tree = tree_unflatten([_asterisk] * treespec.num_leaves, treespec)
     return repr(dummy_tree)
 
 

torch/utils/_pytree.py

@@ -1,3 +1,5 @@
+# Owner(s): ["module: pytree"]
+
 """
 Contains utility functions for working with nested python data structures.
 
@@ -20,6 +22,7 @@ import functools
 import importlib
 import importlib.metadata
 import json
+import sys
 import threading
 import types
 import warnings
@@ -35,23 +38,28 @@ from typing import (
     NoReturn,
     overload,
     Protocol,
+    TYPE_CHECKING,
     TypeAlias,
     TypeVar,
     Union,
 )
-from typing_extensions import deprecated, NamedTuple, Self
+from typing_extensions import deprecated, NamedTuple, Self, TypeIs
 
 from torch.torch_version import TorchVersion as _TorchVersion
 
 
+if TYPE_CHECKING:
+    import torch.utils._cxx_pytree as cxx_pytree
+
+
 __all__ = [
     "PyTree",
     "Context",
     "FlattenFunc",
     "UnflattenFunc",
     "DumpableContext",
-    "ToDumpableContextFn",
-    "FromDumpableContextFn",
+    "ToDumpableContextFunc",
+    "FromDumpableContextFunc",
     "PyTreeSpec",
     "TreeSpec",
     "LeafSpec",
@@ -87,6 +95,9 @@ __all__ = [
 ]
 
 
+__name__ = "torch.utils.pytree.python"  # sets the __module__ attribute of all functions in this module
+
+
 T = TypeVar("T")
 S = TypeVar("S")
 U = TypeVar("U")
@@ -118,17 +129,21 @@ class EnumEncoder(json.JSONEncoder):
         return cast(str, super().default(obj))
 
 
-Context = Any
-PyTree = Any
-FlattenFunc = Callable[[PyTree], tuple[list[Any], Context]]
-UnflattenFunc = Callable[[Iterable[Any], Context], PyTree]
-DumpableContext = Any  # Any json dumpable text
-ToDumpableContextFn = Callable[[Context], DumpableContext]
-FromDumpableContextFn = Callable[[DumpableContext], Context]
-ToStrFunc = Callable[["TreeSpec", list[str]], str]
-MaybeFromStrFunc = Callable[[str], tuple[Any, Context, str] | None]
-KeyPath = tuple[KeyEntry, ...]
-FlattenWithKeysFunc = Callable[[PyTree], tuple[list[tuple[KeyEntry, Any]], Any]]
+Context: TypeAlias = Any
+PyTree: TypeAlias = Any
+FlattenFunc: TypeAlias = Callable[[PyTree], tuple[list[Any], Context]]
+UnflattenFunc: TypeAlias = Callable[[Iterable[Any], Context], PyTree]
+DumpableContext: TypeAlias = Any  # Any json dumpable text
+ToDumpableContextFunc: TypeAlias = Callable[[Context], DumpableContext]
+FromDumpableContextFunc: TypeAlias = Callable[[DumpableContext], Context]
+ToDumpableContextFn: TypeAlias = ToDumpableContextFunc
+FromDumpableContextFn: TypeAlias = FromDumpableContextFunc
+ToStrFunc: TypeAlias = Callable[["TreeSpec", list[str]], str]
+MaybeFromStrFunc: TypeAlias = Callable[[str], tuple[Any, Context, str] | None]
+KeyPath: TypeAlias = tuple[KeyEntry, ...]
+FlattenWithKeysFunc: TypeAlias = Callable[
+    [PyTree], tuple[list[tuple[KeyEntry, Any]], Any]
+]
 
 
 # A NodeDef holds two callables:
@@ -161,8 +176,8 @@ SUPPORTED_NODES: dict[type[Any], NodeDef] = {}
 class _SerializeNodeDef(NamedTuple):
     typ: type[Any]
     serialized_type_name: str
-    to_dumpable_context: ToDumpableContextFn | None
-    from_dumpable_context: FromDumpableContextFn | None
+    to_dumpable_context: ToDumpableContextFunc | None
+    from_dumpable_context: FromDumpableContextFunc | None
 
 
 SUPPORTED_SERIALIZED_TYPES: dict[type[Any], _SerializeNodeDef] = {}
@@ -199,8 +214,8 @@ def register_pytree_node(
     unflatten_fn: UnflattenFunc,
     *,
     serialized_type_name: str | None = None,
-    to_dumpable_context: ToDumpableContextFn | None = None,
-    from_dumpable_context: FromDumpableContextFn | None = None,
+    to_dumpable_context: ToDumpableContextFunc | None = None,
+    from_dumpable_context: FromDumpableContextFunc | None = None,
     flatten_with_keys_fn: FlattenWithKeysFunc | None = None,
 ) -> None:
     """Register a container-like type as pytree node.
@@ -249,9 +264,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,
@@ -527,8 +542,8 @@ def _register_pytree_node(
     maybe_from_str_fn: MaybeFromStrFunc | None = None,  # deprecated
     *,
     serialized_type_name: str | None = None,
-    to_dumpable_context: ToDumpableContextFn | None = None,
-    from_dumpable_context: FromDumpableContextFn | None = None,
+    to_dumpable_context: ToDumpableContextFunc | None = None,
+    from_dumpable_context: FromDumpableContextFunc | None = None,
     flatten_with_keys_fn: FlattenWithKeysFunc | None = None,
 ) -> None:
     """Register a container-like type as pytree node for the Python pytree only.
@@ -594,8 +609,8 @@ def _private_register_pytree_node(
     unflatten_fn: UnflattenFunc,
     *,
     serialized_type_name: str | None = None,
-    to_dumpable_context: ToDumpableContextFn | None = None,
-    from_dumpable_context: FromDumpableContextFn | None = None,
+    to_dumpable_context: ToDumpableContextFunc | None = None,
+    from_dumpable_context: FromDumpableContextFunc | None = None,
     flatten_with_keys_fn: FlattenWithKeysFunc | None = None,
 ) -> None:
     """This is an internal function that is used to register a pytree node type
@@ -1085,7 +1100,9 @@ def _is_leaf(tree: PyTree, is_leaf: Callable[[PyTree], bool] | None = None) -> b
 #   num_children: the number of children of the root Node (i.e., len(children()))
 #   is_leaf(): whether the root Node is a leaf
 @dataclasses.dataclass(init=False, frozen=True, eq=True, repr=False)
-class TreeSpec:
+class PyTreeSpec:
+    """Representing the structure of the pytree."""
+
     type: Any
     _context: Context
     _children: list[Self]
@@ -1164,21 +1181,26 @@ class TreeSpec:
         return self._children
 
     def is_leaf(self) -> bool:
+        """Test whether the treespec represents a leaf."""
         return self.num_nodes == 1 and self.num_leaves == 1
 
     def children(self) -> list[Self]:
+        """Get all the child treespecs."""
         return self._children.copy()
 
     def child(self, index: int) -> Self:
+        """Get the child treespec at the given index."""
         return self._children[index]
 
     def flatten_up_to(self, tree: PyTree) -> list[PyTree]:
-        def helper(treespec: TreeSpec, tree: PyTree, subtrees: list[PyTree]) -> None:
+        """Flatten the subtrees in ``tree`` up to the structure of this treespec and return a list of subtrees."""
+
+        def helper(treespec: TreeSpec, node: PyTree, subtrees: list[PyTree]) -> None:
             if treespec.is_leaf():
-                subtrees.append(tree)
+                subtrees.append(node)
                 return
 
-            node_type = _get_node_type(tree)
+            node_type = _get_node_type(node)
             if treespec.type not in BUILTIN_TYPES:
                 # Always require custom node types to match exactly
                 if node_type != treespec.type:
@@ -1187,7 +1209,7 @@ class TreeSpec:
                         f"expected {treespec.type!r}, but got {node_type!r}.",
                     )
                 flatten_fn = SUPPORTED_NODES[node_type].flatten_fn
-                children, context = flatten_fn(tree)
+                children, context = flatten_fn(node)
                 if len(children) != treespec.num_children:
                     raise ValueError(
                         f"Node arity mismatch; "
@@ -1209,10 +1231,10 @@ class TreeSpec:
                         f"Node type mismatch; "
                         f"expected {treespec.type!r}, but got {node_type!r}.",
                     )
-                if len(tree) != treespec.num_children:
+                if len(node) != treespec.num_children:
                     raise ValueError(
                         f"Node arity mismatch; "
-                        f"expected {treespec.num_children}, but got {len(tree)}.",
+                        f"expected {treespec.num_children}, but got {len(node)}.",
                     )
 
                 if both_standard_dict:
@@ -1224,7 +1246,7 @@ class TreeSpec:
                         else treespec._context[1]
                     )
                     expected_keys = dict_context
-                    got_key_set = set(tree)
+                    got_key_set = set(node)
                     expected_key_set = set(expected_keys)
                     if got_key_set != expected_key_set:
                         missing_keys = expected_key_set.difference(got_key_set)
@@ -1235,11 +1257,11 @@ class TreeSpec:
                         if extra_keys:
                             message += f"; extra key(s): {extra_keys}"
                         raise ValueError(f"Node keys mismatch{message}.")
-                    children = [tree[key] for key in expected_keys]
+                    children = [node[key] for key in expected_keys]
                 else:
                     # node_type is treespec.type
                     flatten_fn = SUPPORTED_NODES[node_type].flatten_fn
-                    children, context = flatten_fn(tree)
+                    children, context = flatten_fn(node)
                     if (
                         node_type is not deque  # ignore mismatch of `maxlen` for deque
                     ) and context != treespec._context:
@@ -1256,6 +1278,7 @@ class TreeSpec:
         return subtrees
 
     def unflatten(self, leaves: Iterable[Any]) -> PyTree:
+        """Reconstruct a pytree from the leaves."""
         if not isinstance(leaves, (list, tuple)):
             leaves = list(leaves)
         if len(leaves) != self.num_leaves:
@@ -1301,7 +1324,7 @@ class TreeSpec:
         return hash((node_type, hashable_context, tuple(self._children)))
 
 
-PyTreeSpec: TypeAlias = TreeSpec
+TreeSpec: TypeAlias = PyTreeSpec
 
 
 # NOTE: subclassing a dataclass is subtle. In order to enable reasoning about
@@ -1363,6 +1386,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: Callable[[PyTree], bool] | None = None,
@@ -1393,11 +1455,14 @@ 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):
-        raise TypeError(
-            f"tree_unflatten(leaves, treespec): Expected `treespec` to be "
-            f"instance of TreeSpec but got item of type {type(treespec)}.",
-        )
+    if not _is_pytreespec_instance(treespec):
+        if not _is_pytreespec_instance(leaves):
+            raise TypeError(
+                f"Expected `treespec` to be an instance of "
+                f"PyTreeSpec but got item of type {type(treespec)}."
+            )
+        # Allow passing the PyTreeSpec instance as the first argument
+        leaves, treespec = treespec, leaves
     return treespec.unflatten(leaves)
 
 
@@ -1827,34 +1892,30 @@ def _broadcast_to_and_flatten(
     treespec: TreeSpec,
     is_leaf: Callable[[PyTree], bool] | None = None,
 ) -> list[Any] | None:
-    if not isinstance(treespec, TreeSpec):
-        raise AssertionError("treespec must be a TreeSpec")
+    def broadcast_prefix(
+        prefix_tree: PyTree,
+        full_tree: PyTree,
+        is_leaf: Callable[[PyTree], bool] | None = 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 +2029,7 @@ _SUPPORTED_PROTOCOLS[1] = _ProtocolFn(_treespec_to_json, _json_to_treespec)
 
 
 def treespec_dumps(treespec: TreeSpec, protocol: int | None = 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
@@ -2001,16 +2058,22 @@ def treespec_loads(serialized: str) -> TreeSpec:
     )
 
 
-class _DummyLeaf:
+class _Asterisk(str):
+    __slots__ = ()
+
+    def __new__(cls) -> Self:
+        return super().__new__(cls, "*")
+
     def __repr__(self) -> str:
-        return "*"
+        return "*"  # no quotes
+
+
+_asterisk = _Asterisk()
+del _Asterisk
 
 
 def treespec_pprint(treespec: TreeSpec) -> str:
-    dummy_tree = tree_unflatten(
-        [_DummyLeaf() for _ in range(treespec.num_leaves)],
-        treespec,
-    )
+    dummy_tree = tree_unflatten([_asterisk] * treespec.num_leaves, treespec)
     return repr(dummy_tree)
 
 

torch/utils/pytree/__init__.py

@@ -0,0 +1,235 @@
+# Owner(s): ["module: pytree"]
+
+"""
+Contains utility functions for working with nested python data structures.
+
+A *pytree* is Python nested data structure. It is a tree in the sense that
+nodes are Python collections (e.g., list, tuple, dict) and the leaves are
+Python values. Furthermore, a pytree should not contain reference cycles.
+
+pytrees are useful for working with nested collections of Tensors. For example,
+one can use `tree_map` to map a function over all Tensors inside some nested
+collection of Tensors and `tree_leaves` to get a flat list of all Tensors
+inside some nested collection. pytrees are helpful for implementing nested
+collection support for PyTorch APIs.
+"""
+
+import os as _os
+import sys as _sys
+from types import ModuleType as _ModuleType
+from typing import Any as _Any, Optional as _Optional, TYPE_CHECKING as _TYPE_CHECKING
+
+import torch.utils._pytree as python
+from torch.utils._pytree import (  # these type aliases are identical in both implementations
+    FlattenFunc,
+    FlattenWithKeysFunc,
+    FromDumpableContextFunc,
+    PyTree,
+    ToDumpableContextFunc,
+    UnflattenFunc,
+)
+
+
+if _TYPE_CHECKING:
+    import torch.utils._cxx_pytree as cxx
+
+
+__all__ = [
+    "PyTreeSpec",
+    "register_pytree_node",
+    "tree_flatten",
+    "tree_unflatten",
+    "tree_iter",
+    "tree_leaves",
+    "tree_structure",
+    "tree_map",
+    "tree_map_",
+    "tree_map_only",
+    "tree_map_only_",
+    "tree_all",
+    "tree_any",
+    "tree_all_only",
+    "tree_any_only",
+    "treespec_pprint",
+    "is_namedtuple",
+    "is_namedtuple_class",
+    "is_namedtuple_instance",
+    "is_structseq",
+    "is_structseq_class",
+    "is_structseq_instance",
+]
+
+
+# NB: Once this variable is read from the environment, the underlying pytree
+#     implementation is frozen. It cannot be swapped to another at runtime.
+PYTORCH_USE_CXX_PYTREE: bool = _os.getenv("PYTORCH_USE_CXX_PYTREE", "0") not in {
+    "0",
+    "",
+}
+
+
+def _import_cxx_pytree_and_store() -> _ModuleType:
+    if not python._cxx_pytree_dynamo_traceable:
+        raise ImportError(
+            "Cannot import package `optree`. "
+            "Please install `optree` via `python -m pip install --upgrade optree`. "
+            "Or set the environment variable `PYTORCH_USE_CXX_PYTREE=0`."
+        )
+
+    import torch.utils._cxx_pytree as cxx
+
+    # This allows the following statements to work properly:
+    #
+    #     import torch.utils.pytree
+    #
+    #     torch.utils.pytree.cxx
+    #     torch.utils.pytree.cxx.tree_map
+    #
+    _sys.modules[f"{__name__}.cxx"] = globals()["cxx"] = cxx
+    return cxx
+
+
+if PYTORCH_USE_CXX_PYTREE:
+    cxx = _import_cxx_pytree_and_store()  # noqa: F811
+else:
+    cxx = _sys.modules.get("torch.utils._cxx_pytree")  # type: ignore[assignment]
+
+
+_sys.modules[f"{__name__}.python"] = python
+if cxx is not None:
+    _sys.modules[f"{__name__}.cxx"] = cxx
+else:
+    del cxx
+
+    class LazyCxxModule(_ModuleType):
+        def __getattr__(self, name: str) -> _Any:
+            if name == "__name__":
+                return f"{__name__}.cxx"
+            if name == "__file__":
+                return python.__file__.removesuffix("_python.py") + "_cxx_pytree.py"
+
+            cxx = globals().get("cxx")
+            if cxx is None:
+                if name.startswith("_"):
+                    raise AttributeError(
+                        f"module {self.__name__!r} has not been imported yet: "
+                        f"accessing attribute {name!r}. "
+                        f"Please import {self.__name__!r} explicitly first."
+                    )
+
+                # Lazy import on first member access
+                cxx = _import_cxx_pytree_and_store()
+
+            return getattr(cxx, name)
+
+        def __setattr__(self, name: str, value: _Any) -> None:
+            # Lazy import
+            cxx = _import_cxx_pytree_and_store()
+            return setattr(cxx, name, value)
+
+    # This allows the following statements to work properly:
+    #
+    #     import torch.utils.pytree.cxx
+    #     from torch.utils.pytree.cxx import tree_map
+    #
+    _sys.modules[f"{__name__}.cxx"] = LazyCxxModule(f"{__name__}.cxx")
+
+    del LazyCxxModule
+
+
+if not PYTORCH_USE_CXX_PYTREE:
+    from torch.utils._pytree import (
+        is_namedtuple,
+        is_namedtuple_class,
+        is_namedtuple_instance,
+        is_structseq,
+        is_structseq_class,
+        is_structseq_instance,
+        PyTreeSpec,
+        register_pytree_node as _register_pytree_node,
+        tree_all,
+        tree_all_only,
+        tree_any,
+        tree_any_only,
+        tree_flatten,
+        tree_iter,
+        tree_leaves,
+        tree_map,
+        tree_map_,
+        tree_map_only,
+        tree_map_only_,
+        tree_structure,
+        tree_unflatten,
+        treespec_pprint,
+    )
+else:
+    from torch.utils._cxx_pytree import (  # type: ignore[assignment,no-redef]
+        is_namedtuple,
+        is_namedtuple_class,
+        is_namedtuple_instance,
+        is_structseq,
+        is_structseq_class,
+        is_structseq_instance,
+        PyTreeSpec,
+        register_pytree_node as _register_pytree_node,
+        tree_all,
+        tree_all_only,
+        tree_any,
+        tree_any_only,
+        tree_flatten,
+        tree_iter,
+        tree_leaves,
+        tree_map,
+        tree_map_,
+        tree_map_only,
+        tree_map_only_,
+        tree_structure,
+        tree_unflatten,
+        treespec_pprint,
+    )
+
+
+def register_pytree_node(
+    cls: type[_Any],
+    /,
+    # intentionally use `*_func` over `*_fn` to match annotations
+    flatten_func: FlattenFunc,
+    unflatten_func: UnflattenFunc,
+) -> None:
+    """Register a container-like type as pytree node.
+
+    Args:
+        cls (type): A Python type to treat as an internal pytree node.
+        flatten_func (callable): A function to be used during flattening, taking an instance of
+            ``cls`` and returning a pair, with (1) an iterable for the children to be flattened
+            recursively, and (2) some hashable auxiliary data to be stored in the treespec and to be
+            passed to the ``unflatten_func``.
+        unflatten_func (callable): A function taking two arguments: the unflattened children, and
+            the auxiliary data that was returned by ``flatten_func`` and stored in the treespec.
+            The function should return an instance of ``cls``.
+
+    Example::
+
+        >>> # xdoctest: +SKIP
+        >>> from collections import UserList
+        ... class MyList(UserList): pass
+        >>> # Registry a Python type with lambda functions
+        ... register_pytree_node(
+        ...     MyList,
+        ...     lambda lst: (list(lst), None),
+        ...     lambda children, _: MyList(children),
+        ... )
+    """
+    _register_pytree_node(
+        cls,
+        flatten_func,
+        unflatten_func,
+    )
+
+
+def __getattr__(name: str) -> _Any:
+    if name == "cxx":
+        # Lazy import
+        return _import_cxx_pytree_and_store()
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

torch/utils/pytree/cxx.pyi

@@ -0,0 +1,8 @@
+# Owner(s): ["module: pytree"]
+
+from .._cxx_pytree import *  # previously public APIs # noqa: F403
+from .._cxx_pytree import (  # non-public internal APIs
+    __all__ as __all__,
+    _broadcast_to_and_flatten as _broadcast_to_and_flatten,
+    KeyPath as KeyPath,
+)

torch/utils/pytree/python.pyi

@@ -0,0 +1,15 @@
+# Owner(s): ["module: pytree"]
+
+from .._pytree import *  # previously public APIs # noqa: F403
+from .._pytree import (  # non-public internal APIs
+    __all__ as __all__,
+    _broadcast_to_and_flatten as _broadcast_to_and_flatten,
+    arg_tree_leaves as arg_tree_leaves,
+    BUILTIN_TYPES as BUILTIN_TYPES,
+    GetAttrKey as GetAttrKey,
+    KeyEntry as KeyEntry,
+    KeyPath as KeyPath,
+    MappingKey as MappingKey,
+    SequenceKey as SequenceKey,
+    SUPPORTED_NODES as SUPPORTED_NODES,
+)