Add correct __all__ for torch.distributed and torch.cuda submodules

test/allowlist_for_publicAPI.json

@@ -141,20 +141,6 @@
   "torch.backends": [
     "contextmanager"
   ],
-  "torch.cpu.amp.autocast_mode": [
-    "Any"
-  ],
-  "torch.cuda": [
-    "Any",
-    "Device",
-    "Dict",
-    "List",
-    "Optional",
-    "Set",
-    "Tuple",
-    "Union",
-    "classproperty"
-  ],
   "torch.cuda.comm": [
     "broadcast",
     "broadcast_coalesced",
@@ -163,12 +149,6 @@
     "scatter",
     "gather"
   ],
-  "torch.cuda.amp.autocast_mode": [
-    "Any"
-  ],
-  "torch.cuda.amp.common": [
-    "find_spec"
-  ],
   "torch.cuda.nccl": [
     "init_rank",
     "is_available",
@@ -292,47 +272,7 @@
   "torch.distributed.optim.utils": [
     "Type"
   ],
-  "torch.distributed.pipeline.sync.checkpoint": [
-    "Checkpoint",
-    "Checkpointing",
-    "Context",
-    "Function",
-    "Recompute",
-    "ThreadLocal",
-    "checkpoint",
-    "enable_checkpointing",
-    "enable_recomputing",
-    "restore_rng_states",
-    "save_rng_states"
-  ],
-  "torch.distributed.pipeline.sync.copy": [
-    "Context",
-    "Copy",
-    "Wait"
-  ],
-  "torch.distributed.pipeline.sync.dependency": [
-    "Fork",
-    "Join",
-    "fork",
-    "join"
-  ],
-  "torch.distributed.pipeline.sync.microbatch": [
-    "Batch",
-    "NoChunk",
-    "check",
-    "gather",
-    "scatter"
-  ],
-  "torch.distributed.pipeline.sync.phony": [
-    "get_phony"
-  ],
   "torch.distributed.pipeline.sync.pipe": [
-    "BalanceError",
-    "PipeSequential",
-    "Pipeline",
-    "WithDevice"
-  ],
-  "torch.distributed.pipeline.sync.pipeline": [
     "Pipeline"
   ],
   "torch.distributed.pipeline.sync.skip.layout": [
@@ -356,25 +296,6 @@
     "current_skip_tracker",
     "use_skip_tracker"
   ],
-  "torch.distributed.pipeline.sync.stream": [
-    "CPUStreamType",
-    "as_cuda",
-    "current_stream",
-    "default_stream",
-    "get_device",
-    "is_cuda",
-    "new_stream",
-    "record_stream",
-    "use_device",
-    "use_stream",
-    "wait_stream"
-  ],
-  "torch.distributed.pipeline.sync.worker": [
-    "Task",
-    "create_workers",
-    "spawn_workers",
-    "worker"
-  ],
   "torch.distributed.remote_device": [
     "Optional",
     "Union"
@@ -395,69 +316,6 @@
     "urlunparse"
   ],
   "torch.distributed.rpc": [
-    "Any",
-    "Dict",
-    "Future",
-    "Generator",
-    "Generic",
-    "GenericWithOneTypeVar",
-    "PyRRef",
-    "RemoteProfilerManager",
-    "RpcAgent",
-    "RpcBackendOptions",
-    "Set",
-    "Store",
-    "TensorPipeAgent",
-    "Tuple",
-    "TypeVar",
-    "WorkerInfo",
-    "enable_gil_profiling",
-    "get_rpc_timeout",
-    "method",
-    "timedelta",
-    "urlparse"
-  ],
-  "torch.distributed.rpc.api": [
-    "Any",
-    "Dict",
-    "Future",
-    "Generic",
-    "GenericWithOneTypeVar",
-    "PyRRef",
-    "PythonUDF",
-    "RPCExecMode",
-    "RemoteProfilerManager",
-    "Set",
-    "TypeVar",
-    "WorkerInfo",
-    "get_rpc_timeout",
-    "method"
-  ],
-  "torch.distributed.rpc.backend_registry": [
-    "Dict",
-    "List",
-    "Set",
-    "Tuple"
-  ],
-  "torch.distributed.rpc.constants": [
-    "timedelta"
-  ],
-  "torch.distributed.rpc.internal": [
-    "Enum"
-  ],
-  "torch.distributed.rpc.options": [
-    "DeviceType",
-    "Dict",
-    "List",
-    "Optional",
-    "Union"
-  ],
-  "torch.distributions.utils": [
-    "Any",
-    "Dict",
-    "Number",
-    "is_tensor_like",
-    "update_wrapper"
   ],
   "torch.fft": [
     "Tensor",

torch/cpu/amp/autocast_mode.py

@@ -1,6 +1,8 @@
 import torch
 from typing import Any
 
+__all__ = ["autocast"]
+
 class autocast(torch.amp.autocast_mode.autocast):
     r"""
     See :class:`torch.autocast`.

torch/cuda/__init__.py

@@ -834,7 +834,7 @@ def dtype(self):
     'IntStorage', 'IntTensor',
     'LongStorage', 'LongTensor',
     'ShortStorage', 'ShortTensor',
-    'CUDAGraph', 'CudaError', 'DeferredCudaCallError', 'Device', 'Event', 'ExternalStream', 'OutOfMemoryError',
+    'CUDAGraph', 'CudaError', 'DeferredCudaCallError', 'Event', 'ExternalStream', 'OutOfMemoryError',
     'Stream', 'StreamContext', 'amp', 'caching_allocator_alloc', 'caching_allocator_delete', 'can_device_access_peer',
     'check_error', 'cudaStatus', 'cudart', 'current_blas_handle', 'current_device', 'current_stream', 'default_generators',
     'default_stream', 'device', 'device_count', 'device_of', 'empty_cache', 'get_arch_list', 'get_device_capability',

torch/cuda/amp/autocast_mode.py

@@ -9,6 +9,7 @@
 from torch._six import string_classes
 from typing import Any
 
+__all__ = ["autocast", "custom_fwd", "custom_bwd"]
 
 class autocast(torch.amp.autocast_mode.autocast):
     r"""

torch/cuda/amp/common.py

@@ -1,6 +1,7 @@
 import torch
 from importlib.util import find_spec
 
+__all__ = ["amp_definitely_not_available"]
 
 def amp_definitely_not_available():
     return not (torch.cuda.is_available() or find_spec('torch_xla'))

torch/distributed/__init__.py

@@ -4,7 +4,6 @@
 
 import torch
 
-
 def is_available() -> bool:
     """
     Returns ``True`` if the distributed package is available. Otherwise,

torch/distributed/pipeline/sync/checkpoint.py

@@ -47,7 +47,9 @@
 from .microbatch import Batch
 from .phony import get_phony
 
-__all__ = ["is_checkpointing", "is_recomputing"]
+__all__ = ["Function", "checkpoint", "Checkpointing", "ThreadLocal", "enable_checkpointing",
+           "enable_recomputing", "is_checkpointing", "is_recomputing", "Context", "save_rng_states",
+           "restore_rng_states", "Checkpoint", "Recompute"]
 
 
 Tensors = Sequence[Tensor]

torch/distributed/pipeline/sync/copy.py

@@ -15,7 +15,7 @@
 
 from .stream import AbstractStream, current_stream, get_device, record_stream, use_stream, wait_stream
 
-__all__: List[str] = []
+__all__: List[str] = ["Context", "Copy", "Wait"]
 
 
 Tensors = Sequence[Tensor]

torch/distributed/pipeline/sync/dependency.py

@@ -12,7 +12,7 @@
 
 from .phony import get_phony
 
-__all__: List[str] = []
+__all__: List[str] = ["fork", "Fork", "join", "Join"]
 
 
 def fork(input: Tensor) -> Tuple[Tensor, Tensor]:

torch/distributed/pipeline/sync/microbatch.py

@@ -12,7 +12,7 @@
 from torch import Tensor
 import torch.cuda.comm
 
-__all__: List[str] = []
+__all__: List[str] = ["NoChunk", "Batch", "check", "scatter", "gather"]
 
 
 Tensors = Sequence[Tensor]

torch/distributed/pipeline/sync/phony.py

@@ -12,7 +12,7 @@
 
 from .stream import default_stream, use_stream
 
-__all__: List[str] = []
+__all__: List[str] = ["get_phony"]
 
 
 _phonies: Dict[Tuple[torch.device, bool], Tensor] = {}

torch/distributed/pipeline/sync/pipe.py

@@ -21,7 +21,7 @@
 from .skip.skippable import verify_skippables
 from .stream import AbstractStream, new_stream
 
-__all__ = ["Pipe"]
+__all__ = ["Pipe", "BalanceError", "PipeSequential", "WithDevice"]
 
 
 Device = Union[torch.device, int, str]

torch/distributed/pipeline/sync/pipeline.py

@@ -23,7 +23,7 @@
 from .stream import AbstractStream, current_stream, use_device
 from .worker import Task, create_workers
 
-__all__: List[str] = []
+__all__: List[str] = ["Pipeline"]
 
 
 Tensors = Sequence[Tensor]

torch/distributed/pipeline/sync/stream.py

@@ -12,7 +12,9 @@
 
 import torch
 
-__all__: List[str] = []
+__all__: List[str] = ["CPUStreamType", "new_stream", "current_stream", "default_stream",
+                      "use_device", "use_stream", "get_device", "wait_stream", "record_stream",
+                      "is_cuda", "as_cuda"]
 
 
 class CPUStreamType:

torch/distributed/pipeline/sync/utils.py

@@ -1,6 +1,8 @@
 from torch import nn
 from typing import List
 
+__all__ = ["partition_model"]
+
 def partition_model(
         module: nn.Sequential,
         balance: List[int],

torch/distributed/pipeline/sync/worker.py

@@ -17,7 +17,7 @@
 from .microbatch import Batch
 from .stream import AbstractStream, use_device, use_stream
 
-__all__: List[str] = []
+__all__: List[str] = ["Task", "worker", "create_workers", "spawn_workers"]
 
 
 ExcInfo = Tuple[Type[BaseException], BaseException, TracebackType]

torch/distributed/rpc/__init__.py

@@ -9,13 +9,13 @@
 import torch
 import torch.distributed as dist
 
-
 logger = logging.getLogger(__name__)
 
 
 _init_counter = 0
 _init_counter_lock = threading.Lock()
 
+__all__ = ["is_available"]
 
 def is_available():
     return hasattr(torch._C, "_rpc_init")
@@ -77,6 +77,9 @@ def is_available():
 
     rendezvous_iterator: Generator[Tuple[Store, int, int], None, None]
 
+    __all__ += ["init_rpc", "BackendType", "TensorPipeRpcBackendOptions"]
+    __all__ = __all__ + api.__all__ + backend_registry.__all__
+
     def init_rpc(
         name,
         backend=None,

torch/distributed/rpc/backend_registry.py

@@ -11,6 +11,9 @@
 from . import api
 from . import constants as rpc_constants
 
+__all__ = ["backend_registered", "register_backend", "construct_rpc_backend_options", "init_backend",
+           "BackendValue", "BackendType"]
+
 BackendValue = collections.namedtuple(
     "BackendValue", ["construct_rpc_backend_options_handler", "init_backend_handler"]
 )

torch/distributed/rpc/constants.py

@@ -1,5 +1,5 @@
 from datetime import timedelta
-
+from typing import List
 from torch._C._distributed_rpc import (
     _DEFAULT_INIT_METHOD,
     _DEFAULT_NUM_WORKER_THREADS,
@@ -20,3 +20,5 @@
 DEFAULT_PROCESS_GROUP_TIMEOUT: timedelta = timedelta(milliseconds=2 ** 31 - 1)
 # Value indicating that timeout is not set for RPC call, and the default should be used.
 UNSET_RPC_TIMEOUT: float = _UNSET_RPC_TIMEOUT
+
+__all__: List[str] = []

torch/distributed/rpc/internal.py

@@ -11,6 +11,7 @@
 import torch.distributed as dist
 from torch._C._distributed_rpc import _get_current_rpc_agent
 
+__all__ = ["RPCExecMode", "serialize", "deserialize", "PythonUDF", "RemoteException"]
 
 # Thread local tensor tables to store tensors while pickling torch.Tensor
 # objects

torch/distributed/rpc/options.py

@@ -7,6 +7,7 @@
 
 DeviceType = Union[int, str, torch.device]
 
+__all__ = ["TensorPipeRpcBackendOptions"]
 
 def _to_device(device: DeviceType) -> torch.device:
     device = torch.device(device)

torch/distributions/utils.py

@@ -7,6 +7,8 @@
 
 euler_constant = 0.57721566490153286060  # Euler Mascheroni Constant
 
+__all__ = ["broadcast_all", "logits_to_probs", "clamp_probs", "probs_to_logits", "lazy_property",
+           "tril_matrix_to_vec", "vec_to_tril_matrix"]
 
 def broadcast_all(*values):
     r"""