pytorch · XuehaiPan · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024
diff --git a/.lintrunner.toml b/.lintrunner.toml
@@ -177,6 +177,7 @@ include_patterns = [
     'benchmarks/instruction_counts/**/*.py',
     'tools/**/*.py',
     'torchgen/**/*.py',
+    'torch/utils/pytree/__init__.py',
     'torch/utils/_pytree.py',
     'torch/utils/_cxx_pytree.py',
     'torch/utils/benchmark/utils/common.py',

diff --git a/docs/source/pytorch-api.md b/docs/source/pytorch-api.md
@@ -67,6 +67,7 @@ sparse
 storage
 torch.testing <testing>
 torch.utils <utils>
+torch.utils.pytree
 torch.utils.benchmark <benchmark_utils>
 torch.utils.bottleneck <bottleneck>
 torch.utils.checkpoint <checkpoint>

diff --git a/docs/source/torch.utils.pytree.rst b/docs/source/torch.utils.pytree.rst
@@ -0,0 +1,7 @@
+torch.utils.pytree
+==================
+
+.. currentmodule:: torch.utils.pytree
+
+.. automodule:: torch.utils.pytree
+    :members:
diff --git a/mypy-strict.ini b/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,

diff --git a/test/dynamo/test_flat_apply.py b/test/dynamo/test_flat_apply.py
@@ -147,8 +147,8 @@ def forward(self, L_x_: "f32[10]", L_y_: "f32[10]"):
 
         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.PyTreeSpec = self.trace_point_tensor_spec
+        trace_point_tensor_input_spec : torch.utils.pytree.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