From ac7fd643ff2a5a687289e561d1ccf36eaeaee24c Mon Sep 17 00:00:00 2001
From: Xuehai Pan <XuehaiPan@pku.edu.cn>
Date: Fri, 10 Jan 2025 01:12:17 +0800
Subject: [PATCH] Update (base update)

[ghstack-poisoned]
---
 .lintrunner.toml               |  2 +-
 mypy-strict.ini                |  2 +-
 torch/utils/pytree/__init__.py | 15 ++++++++++-----
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/.lintrunner.toml b/.lintrunner.toml
index 33602f74ef8..1aa399af61e 100644
--- a/.lintrunner.toml
+++ b/.lintrunner.toml
@@ -171,7 +171,7 @@ include_patterns = [
     'benchmarks/instruction_counts/**/*.py',
     'tools/**/*.py',
     'torchgen/**/*.py',
-    'torch/utils/pytree.py',
+    'torch/utils/pytree/__init__.py',
     'torch/utils/_pytree.py',
     'torch/utils/_cxx_pytree.py',
     'torch/utils/benchmark/utils/common.py',
diff --git a/mypy-strict.ini b/mypy-strict.ini
index a2815f05fb2..2e92fad971f 100644
--- a/mypy-strict.ini
+++ b/mypy-strict.ini
@@ -29,7 +29,7 @@ files =
     benchmarks/instruction_counts,
     tools,
     torch/profiler/_memory_profiler.py,
-    torch/utils/pytree.py,
+    torch/utils/pytree/__init__.py,
     torch/utils/_pytree.py,
     torch/utils/_cxx_pytree.py,
     torch/utils/benchmark/utils/common.py,
diff --git a/torch/utils/pytree/__init__.py b/torch/utils/pytree/__init__.py
index d5b44c455e9..a3daff0e8f0 100644
--- a/torch/utils/pytree/__init__.py
+++ b/torch/utils/pytree/__init__.py
@@ -77,14 +77,19 @@ __all__ = [
     "tree_all_only",
     "tree_any_only",
     "treespec_pprint",
-    "FlattenFunc",
-    "UnflattenFunc",
-    "FlattenWithKeysFunc",
-    "ToDumpableContextFunc",
-    "FromDumpableContextFunc",
 ]
 
 
+if _TYPE_CHECKING:
+    __all__ += [
+        "FlattenFunc",
+        "UnflattenFunc",
+        "FlattenWithKeysFunc",
+        "ToDumpableContextFunc",
+        "FromDumpableContextFunc",
+    ]
+
+
 PYTORCH_USE_CXX_PYTREE: bool = _os.getenv("PYTORCH_USE_CXX_PYTREE", "0") not in {
     "0",
     "",