Name	Size	Mode	Actions
__pycache__/	-	0755	rm
comm.py	10788	0644	editdlrm
common_types.pyi	123	0644	editdlrm
data_parallel.py	10270	0644	editdlrm
data_parallel.pyi	717	0644	editdlrm
distributed.py	69172	0644	editdlrm
distributed.pyi	860	0644	editdlrm
parallel_apply.py	3216	0644	editdlrm
parallel_apply.pyi	293	0644	editdlrm
replicate.py	6007	0644	editdlrm
replicate.pyi	245	0644	editdlrm
scatter_gather.py	3136	0644	editdlrm
scatter_gather.pyi	989	0644	editdlrm
_functions.py	4664	0644	editdlrm
__init__.py	645	0644	editdlrm
__init__.pyi	336	0644	editdlrm

from typing import Any, Dict, List, Tuple, overload, TypeVar from ... import Tensor from .common_types import _device_t, _devices_t T = TypeVar('T', Dict, List, Tuple) # For some reason, 'scatter' returns a tuple when given a single Tensor input but a list otherwise. @overload def scatter(inputs: Tensor, target_gpus: _devices_t, dim: int = ...) -> Tuple[Tensor, ...]: ... # flake8 will raise a spurious error here since `torch/__init__.pyi` has not been generated yet # so mypy will interpret `Tensor` as `Any` since it is an import from what it believes to be an # untyped module. Thus to mypy, the first definition of `scatter` looks strictly more general # than this overload. @overload def scatter(inputs: T, target_gpus: _devices_t, dim: int = ...) -> List[T]: ... # TODO More precise types here. def scatter_kwargs(inputs: Any, kwargs: Any, target_gpus: _devices_t, dim: int = ...) -> Any: ... def gather(outputs: Any, target_device: _device_t, dim: int = ...) -> Any: ...