Add __all__ to torch.utils submodules (#85331)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/85331
Approved by: https://github.com/albanD

test/allowlist_for_publicAPI.json

@@ -1373,10 +1373,6 @@
   "torch.random": [
     "Generator"
   ],
-  "torch.return_types": [
-    "attr",
-    "pytree_register_structseq"
-  ],
   "torch.serialization": [
     "Any",
     "BinaryIO",
@@ -1550,18 +1546,6 @@
     "Serialization",
     "wrapper_singleton"
   ],
-  "torch.utils.cpp_extension": [
-    "ExtensionVersioner",
-    "FileBaton",
-    "GeneratedFileCleaner",
-    "List",
-    "Optional",
-    "TorchVersion",
-    "Tuple",
-    "Union",
-    "build_ext",
-    "get_hip_file_path"
-  ],
   "torch.utils.data": [
     "_DatasetKind",
     "argument_validation",
@@ -1586,26 +1570,6 @@
     "Any",
     "to_dlpack"
   ],
-  "torch.utils.hipify.hipify_python": [
-    "Dict",
-    "HipifyFinalResult",
-    "HipifyResult",
-    "Iterable",
-    "Iterator",
-    "List",
-    "Mapping",
-    "Optional"
-  ],
-  "torch.utils.hooks": [
-    "Any",
-    "OrderedDict"
-  ],
-  "torch.utils.show_pickle": [
-    "Any",
-    "BinaryIO",
-    "IO",
-    "Union"
-  ],
   "torch": [
     "BFloat16Storage",
     "BFloat16Tensor",

torch/return_types.py

@@ -1,7 +1,7 @@
 import torch
 import inspect
 
-__all__ = []
+__all__ = ["pytree_register_structseq"]
 
 # error: Module has no attribute "_return_types"
 return_types = torch._C._return_types  # type: ignore[attr-defined]
@@ -19,8 +19,8 @@ for name in dir(return_types):
     if name.startswith('__'):
         continue
 
-    attr = getattr(return_types, name)
-    globals()[name] = attr
+    _attr = getattr(return_types, name)
+    globals()[name] = _attr
 
     if not name.startswith('_'):
         __all__.append(name)
@@ -30,5 +30,5 @@ for name in dir(return_types):
     # is no longer the case.
     # NB: I don't know how to check that something is a "structseq" so we do a fuzzy
     # check for tuple
-    if inspect.isclass(attr) and issubclass(attr, tuple):
-        pytree_register_structseq(attr)
+    if inspect.isclass(_attr) and issubclass(_attr, tuple):
+        pytree_register_structseq(_attr)

torch/utils/cpp_extension.py

@@ -70,7 +70,9 @@ CUDA_CLANG_VERSIONS = {
     '11.7': ((6, 0, 0), (13, 0, 0)),
 }
 
-
+__all__ = ["get_default_build_root", "check_compiler_ok_for_platform", "get_compiler_abi_compatibility_and_version", "BuildExtension",
+           "CppExtension", "CUDAExtension", "include_paths", "library_paths", "load", "load_inline", "is_ninja_available",
+           "verify_ninja_availability"]
 # Taken directly from python stdlib < 3.9
 # See https://github.com/pytorch/pytorch/issues/48617
 def _nt_quote_args(args: Optional[List[str]]) -> List[str]:

torch/utils/hooks.py

@@ -4,6 +4,7 @@ import weakref
 import warnings
 from typing import Any
 
+__all__ = ["RemovableHandle", "unserializable_hook", "warn_if_has_hooks", "BackwardHook"]
 
 class RemovableHandle(object):
     """A handle which provides the capability to remove a hook."""

torch/utils/show_pickle.py

@@ -7,6 +7,7 @@ import zipfile
 import fnmatch
 from typing import Any, IO, BinaryIO, Union
 
+__all__ = ["FakeObject", "FakeClass", "DumpUnpickler", "main"]
 
 class FakeObject(object):
     def __init__(self, module, name, args):