__init__.pyi.in

# ${generated_comment}

import torch
from torch import Tensor
from enum import Enum
from pathlib import Path
from typing import (Any, BinaryIO, Callable, ContextManager, Dict, Iterator, List, NamedTuple,
                    Optional, overload, Sequence, Tuple, TypeVar, Type, Union, Generic,
                    Set, AnyStr)
from torch._six import inf

from torch.types import _int, _float, _bool, _dtype, _device, _qscheme, _size, _layout, Device, Number, Storage

import builtins

# This module is defined in torch/csrc/Module.cpp

from . import _nn as _nn
from . import _onnx as _onnx
from . import _VariableFunctions as _VariableFunctions

T = TypeVar('T')

# Defined in torch/csrc/Device.cpp
class device:
    type: str  # THPDevice_type
    index: _int  # THPDevice_index

    # THPDevice_pynew
    @overload
    def __init__(self, device: Union[_device, _int, str]) -> None: ...

    @overload
    def __init__(self, type: str, index: _int) -> None: ...

    def __reduce__(self) -> Tuple[Any, ...]: ...  # THPDevice_reduce

# Defined in torch/csrc/Stream.cpp
class Stream:
    _cdata: _int  # Stream handle
    device: device # The device of the stream

    ...

# Defined in torch/csrc/Size.cpp
class Size(Tuple[_int, ...]):
    # TODO: __reduce__

    @overload
    def __getitem__(self: Size, key: _int) -> _int: ...

    @overload
    def __getitem__(self: Size, key: slice) -> Size: ...

    def numel(self: Size) -> _int: ...

    ...

# Defined in torch/csrc/Dtype.cpp
class dtype:
    # TODO: __reduce__
    is_floating_point: _bool
    is_complex: _bool
    is_signed: _bool
    ...

# Defined in torch/csrc/TypeInfo.cpp
class iinfo:
    bits: _int
    min: _int
    max: _int
    dtype: str

    def __init__(self, dtype: _dtype) -> None: ...

class finfo:
    bits: _float
    min: _float
    max: _float
    eps: _float
    tiny: _float
    resolution: _float
    dtype: str

    @overload
    def __init__(self, dtype: _dtype) -> None: ...

    @overload
    def __init__(self) -> None: ...

${dtype_class_hints}

# Defined in torch/csrc/Layout.cpp
class layout:
    ...

# Defined in torch/csrc/utils/disable_torch_function.cpp
def DisableTorchFunction(): ...

# Defined in torch/csrc/utils/tensor_layouts.cpp
strided : layout = ...
sparse_coo : layout = ...
_mkldnn : layout = ...

# Defined in torch/csrc/MemoryFormat.cpp
class memory_format: ...

# Defined in torch/csrc/utils/tensor_memoryformats.cpp
contiguous_format: memory_format = ...
channels_last: memory_format = ...
channels_last_3d: memory_format = ...
preserve_format: memory_format = ...

# Defined in torch/csrc/QScheme.cpp
class qscheme: ...

# Defined in torch/csrc/utils/tensor_qschemes.h
per_tensor_affine: qscheme = ...
per_channel_affine: qscheme = ...
per_tensor_symmetric: qscheme = ...
per_channel_symmetric: qscheme = ...
per_channel_affine_float_qparams: qscheme = ...

# Defined in torch/csrc/autograd/python_function.cpp
class _FunctionBase(object):
    ...

# Defined in torch/csrc/autograd/python_legacy_variable.cpp
class _LegacyVariableBase(object):
    def __init__(
        self,
        data: Optional[Tensor]=...,
        requires_grad: Optional[_bool]=...,
        volatile: Optional[_bool]=...,
        _grad_fn: Optional[_FunctionBase]=...
    ) -> None: ...

# Defined in torch/csrc/jit/python/init.cpp
class IODescriptor: ...

class JITException: ...

class Future(object):
  def __init__(self) -> None: ...
  def done(self) -> _bool: ...
  def wait(self) -> Any: ...
  def add_done_callback(self, callback: Callable) -> None: ...
  def then(self, callback: Callable) -> Future: ...
  def set_result(self, result: Any) -> None: ...

def _jit_set_num_profiled_runs(num: _size) -> _size: ...

# Defined in torch/csrc/jit/passes/xnnpack_rewrite.h
class MobileOptimizerType:
    ...

CONV_BN_FUSION: MobileOptimizerType
INSERT_FOLD_PREPACK_OPS: MobileOptimizerType
REMOVE_DROPOUT: MobileOptimizerType
FUSE_ADD_RELU: MobileOptimizerType
HOIST_CONV_PACKED_PARAMS: MobileOptimizerType

def fork(*args: Any, **kwargs: Any) -> Future: ...
def wait(fut: Future) -> Any: ...
def _collect_all(futures: List[Future]) -> Future: ...

def unify_type_list(types: List[JitType]) -> JitType: ...
def _freeze_module(module: ScriptModule, preserved_attrs: List[str], freeze_interfaces: _bool = True) -> ScriptModule: ...
def _is_tracing() -> _bool: ...
def _jit_init() -> _bool: ...
def _jit_flatten(arg: Any) -> Tuple[List[Tensor], IODescriptor]: ...
def _jit_unflatten(vars: List[Tensor], desc: IODescriptor) -> Any: ...
def _jit_get_operation(op_name: str) -> Callable: ...
def _jit_pass_optimize_for_mobile(module: 'torch.jit.ScriptModule',
                                  optimization_blocklist: Set[MobileOptimizerType],
                                  preserved_methods: List[AnyStr]) -> 'torch.jit.ScriptModule': ...
def _jit_pass_vulkan_optimize_for_mobile(module: 'torch.jit.ScriptModule',
                                         preserved_methods: List[AnyStr]) -> 'torch.jit.ScriptModule': ...
def _jit_pass_metal_optimize_for_mobile(module: 'torch.jit.ScriptModule',
                                         preserved_methods: List[AnyStr]) -> 'torch.jit.ScriptModule': ...
def _jit_pass_inline(Graph) -> None: ...
def _jit_get_schemas_for_operator(name :str) -> List[FunctionSchema]: ...
def _jit_can_fuse_on_cpu() -> _bool: ...
def _jit_can_fuse_on_gpu() -> _bool: ...
def _jit_texpr_fuser_enabled() -> _bool: ...
def _jit_nvfuser_enabled() -> _bool: ...
def _jit_override_can_fuse_on_cpu(override: _bool): ...
def _jit_override_can_fuse_on_gpu(override: _bool): ...
def _jit_set_texpr_fuser_enabled(enable: _bool): ...
def _jit_set_nvfuser_enabled(enable: _bool) -> _bool: ...
def _jit_pass_canonicalize(graph: Graph): ...
def _jit_pass_erase_shape_information(graph: Graph): ...
def _jit_pass_fold_convbn(module: 'torch.jit.ScriptModule'): ...
def _jit_pass_insert_observers(module: 'torch.jit.ScriptModule',
                               method_name: str,
                               qconfig_dict: Dict[str, Any],
                               inplace: _bool,
                               quant_type: _int): ...
def _jit_pass_insert_quant_dequant(module: 'torch.jit.ScriptModule',
                                   method_name: str,
                                   inplace: _bool,
                                   debug: _bool,
                                   quant_type: _int): ...
def _jit_pass_quant_finalize(module: 'torch.jit.ScriptModule',
                             quant_type: _int,
                             preserved_attrs: Sequence[str]): ...
def _jit_set_profiling_executor(profiling_flag: _bool) -> _bool: ...
def _jit_set_profiling_mode(profiling_flag: _bool) -> _bool: ...
def _jit_try_infer_type(obj: Any) -> JitType: ...
def _jit_get_trigger_value(trigger_name: str) -> _int: ...

# Defined in torch/csrc/jit/python/script_init.cpp
ResolutionCallback = Callable[[str], Callable[..., Any]]

def _create_function_from_graph(qualname: str, graph: Graph) -> Graph: ...
def _debug_set_autodiff_subgraph_inlining(disabled: _bool) -> None: ...
def _ivalue_tags_match(lhs: ScriptModule, rhs: ScriptModule) -> _bool: ...
def _jit_clear_class_registry() -> None: ...
def _jit_set_emit_hooks(ModuleHook: Optional[Callable], FunctionHook: Optional[Callable]) -> None: ...
def _jit_get_emit_hooks() -> Tuple[Callable, Callable]: ...
def _load_for_lite_interpreter(filename: Union[str, Path], map_location: Union[_device, str, None]): ...
def _load_for_lite_interpreter_from_buffer(buffer: BinaryIO, map_location: Union[_device, str, None]): ...
def _logging_set_logger(logger: LoggerBase) -> LoggerBase: ...
def _get_graph_executor_optimize() -> _bool: ...
def _set_graph_executor_optimize(optimize: _bool): ...
def _export_opnames(module: ScriptModule) -> List[str]: ...
def _create_function_from_trace(
    qualname: str,
    func: Callable[..., Any],
    input_tuple: Tuple[Any, ...],
    var_lookup_fn: Callable[[Tensor], str],
    strict: _bool,
    force_outplace: _bool
) -> Tuple[Graph, Stack]: ...
def _jit_is_script_object(obj: Any) -> _bool: ...
def _last_executed_optimized_graph() -> Graph: ...
def parse_type_comment(comment: str) -> Decl: ...
def merge_type_from_type_comment(decl: Decl, type_annotation_decl: Decl, is_method: _bool) -> Decl: ...
def _resolve_type_from_object(obj: Any, range: SourceRange, rcb: ResolutionCallback) -> JitType: ...
def _create_module_with_type(ty: JitType) -> ScriptModule: ...
def _run_emit_module_hook(m: ScriptModule): ...
def _replace_overloaded_method_decl(overload_decl: Decl, implementation_def: Def, new_name: str) -> Def: ...
def _jit_script_interface_compile(name: str, class_def: ClassDef, rcb: ResolutionCallback, is_module: _bool): ...
def _jit_script_compile_overload(
    qualname: str,
    overload_decl: Decl,
    implementation_def: Def,
    rcb: ResolutionCallback,
    implementation_defaults: Dict[str, Any],
    signature: Any
): ...
def _jit_script_compile(
    qual_name: str,
    definition: Def,
    rcb: ResolutionCallback,
    defaults: Dict[str, Any]
): ...
def _jit_script_class_compile(
    qual_name: str,
    definition: ClassDef,
    defaults: Dict[str, Dict[str, Any]],
    rcb: ResolutionCallback,
): ...
def _parse_source_def(src: str) -> Def: ...
def import_ir_module(
    cu: CompilationUnit,
    filename: Union[str, Path],
    map_location: Union[_device, str, None],
    extra_files: Dict[str, Any]
) -> ScriptModule: ...
def import_ir_module_from_buffer(
    cu: CompilationUnit,
    buffer: BinaryIO,
    map_location: Union[_device, str, None],
    extra_files: Dict[str, Any]
) -> ScriptModule: ...

# Defined in torch/torch/csrc/jit/ir/ir.h
class Graph:
    ...

# Defined in torch/csrc/jit/ir/ir.h
class Value:
    ...

# Defined in torch/csrc/jit/ir/ir.h
class Block:
    ...

# Defined in torch/csrc/jit/ir/ir.h
class Node:
    ...


# Defined in torch/aten/src/ATen/core/function_schema.h
class Argument:
    name: str
    type: JitType
    default_value: Optional[Any]
    def has_default_value(self) -> _bool: ...
    ...
class FunctionSchema:
    arguments: List[Argument]
    returns: List[Argument]
    ...

# Defined in torch/csrc/jit/python/script_init.cpp
class ConcreteModuleTypeBuilder:
    def __init__(self, obj: Any) -> None: ...
    def set_module_dict(self): ...
    def set_module_list(self): ...
    def add_attribute(self, name: str, ty: JitType, is_param: _bool, is_buffer: _bool): ...
    def add_module(self, name: str, meta: ConcreteModuleType): ...
    def add_constant(self, name: str, value: Any): ...
    def add_overload(self, method_name: str, overloaded_method_names: List[str]): ...
    def add_builtin_function(self, name: str, symbol_name: str): ...
    def add_failed_attribute(self, name: str, failure_reason: str): ...
    def add_function_attribute(self, name: str, ty: JitType, func: Callable[..., Any]): ...
    def add_ignored_attribute(self, name: str): ...
    def add_ignored_attributes(self, names: List[str]): ...

class ConcreteModuleType:
    def get_constants(self) -> Dict[str, Any]: ...
    def equals(self, other: 'ConcreteModuleType') -> _bool: ...

    @staticmethod
    def from_jit_type(ty: JitType) -> ConcreteModuleType: ...

class CallStack:
    def __init__(self, name: str, range: SourceRange): ...

class ErrorReport:
    def __init__(self, range: SourceRange) -> None: ...
    def what(self) -> str: ...

    @staticmethod
    def call_stack() -> str: ...

class CompilationUnit:
    def __init__(self) -> None: ...
    def find_function(self, name: str) -> ScriptFunction: ...
    def define(self, script: str, rcb: ResolutionCallback): ...
    def get_interface(self, name: str) -> InterfaceType: ...

class ScriptModule:
    def setattr(self, name: str, value: Any): ...
    def _method_names(self) -> List[str]: ...
    def _get_method(self, name: str) -> ScriptMethod: ...

class ScriptFunction:
    def __call__(self, *args, **kwargs) -> Tensor: ...
    def save(self, filename: str, _extra_files: Dict[str, bytes]) -> None: ...
    def save_to_buffer(self, _extra_files: Dict[str, bytes]) -> bytes: ...
    def graph(self) -> Graph: ...
    def inlined_graph(self) -> Graph: ...
    def schema(self) -> FunctionSchema: ...
    def code(self) -> str: ...
    def name(self) -> str: ...
    @property
    def qualified_name(self) -> str: ...

class ScriptMethod:
    ...

class ModuleDict:
    def __init__(self, mod: ScriptModule) -> None: ...
    def items(self) -> List[Tuple[str, Any]]: ...

class ParameterDict:
    def __init__(self, mod: ScriptModule) -> None: ...

class BufferDict:
    def __init__(self, mod: ScriptModule) -> None: ...

# Defined in torch/csrc/Module.cpp
def _initExtension(shm_manager_path: str) -> None: ...  # THPModule_initExtension
def _autograd_init() -> _bool: ...  # THPAutograd_initExtension
def _add_docstr(obj: T, doc_obj: str) -> T: ...  # THPModule_addDocStr
def _init_names(arg: Sequence[Type]) -> None: ...  # THPModule_initNames
def _has_distributed() -> _bool: ...  # THPModule_hasDistributed
def _set_default_tensor_type(type) -> None: ...  # THPModule_setDefaultTensorType
def _set_default_dtype(d: _dtype) -> None: ...  # THPModule_setDefaultDtype
def _infer_size(arg1: Size, arg2: Size) -> Size: ...  # THPModule_inferSize
def _crash_if_csrc_asan() -> _int: ...  # THPModule_crashIfCsrcASAN
def _crash_if_csrc_ubsan() -> _int: ...  # THPModule_crashIfCsrcUBSAN
def _crash_if_aten_asan() -> _int: ...  # THPModule_crashIfATenASAN
def _show_config() -> str: ...  # THPModule_showConfig
def _parallel_info() -> str: ...  # THPModule_parallelInfo
def _set_backcompat_broadcast_warn(arg: _bool) -> None: ...  # THPModule_setBackcompatBroadcastWarn
def _get_backcompat_broadcast_warn() -> _bool: ...  # THPModule_getBackcompatBroadcastWarn
def _set_backcompat_keepdim_warn(arg: _bool) -> None: ...  # THPModule_setBackcompatKeepdimWarn
def _get_backcompat_keepdim_warn() -> _bool: ...  # THPModule_getBackcompatKeepdimWarn
def get_num_thread() -> _int: ...  # THPModule_getNumThreads
def set_num_threads(nthreads: _int) -> None: ...  # THPModule_setNumThreads
def get_num_interop_threads() -> _int: ...  # THPModule_getNumInteropThreads
def set_num_interop_threads(nthreads: _int) -> None: ...  # THPModule_setNumInteropThreads
def _get_cudnn_enabled() -> _bool: ...  # THPModule_userEnabledCuDNN
def _set_cudnn_enabled(arg: _bool) -> None: ...  # THPModule_setUserEnabledCuDNN
def _get_mkldnn_enabled() -> _bool: ...  # THPModule_userEnabledMkldnn
def _set_mkldnn_enabled(arg: _bool) -> None: ...  # THPModule_setUserEnabledMkldnn
def _get_cudnn_benchmark() -> _bool: ...  # THPModule_benchmarkCuDNN
def _set_cudnn_benchmark(arg: _bool) -> None: ...  # THPModule_setBenchmarkCuDNN
def _get_cudnn_deterministic() -> _bool: ...  # THPModule_deterministicCuDNN
def _set_cudnn_deterministic(arg: _bool) -> None: ...  # THPModule_setDeterministicCuDNN
def _get_deterministic() -> _bool: ...  # THPModule_deterministic
def _set_deterministic(arg: _bool) -> None: ...  # THPModule_setDeterministic
def _get_cudnn_allow_tf32() -> _bool: ...  # THPModule_allowTF32CuDNN
def _set_cudnn_allow_tf32(arg: _bool) -> None: ...  # THPModule_setAllowTF32CuDNN
def _get_cublas_allow_tf32() -> _bool: ...  # THPModule_allowTF32CuBLAS
def _set_cublas_allow_tf32(arg: _bool) -> None: ...  # THPModule_setAllowTF32CuBLAS
# NB: There is no Capsule type in typing, see
# https://code.activestate.com/lists/python-dev/139675/
def _to_dlpack(data: Tensor) -> Any: ...  # THPModule_toDLPack
def _from_dlpack(data: Any) -> Tensor: ...  # THPModule_fromDLPack
def set_flush_denormal(arg: _bool) -> _bool: ...  # THPModule_setFlushDenormal
def get_default_dtype() -> _dtype: ...  # THPModule_getDefaultDtype
def _get_default_device() -> str: ...  # THPModule_getDefaultDevice
def _get_qengine() -> _int: ...  # THPModule_qEngine
def _set_qengine(qegine: _int) -> None: ...  # THPModule_setQEngine
def _supported_qengines() -> List[_int]: ...  # THPModule_supportedQEngines
def _is_xnnpack_enabled() -> _bool: ...  # THPModule_isEnabledXNNPACK
def _vmapmode_increment_nesting() -> _int: ...  # THPModule_vmapmode_increment_nesting
def _vmapmode_decrement_nesting() -> _int: ...  # THPModule_vmapmode_decrement_nesting
def _log_api_usage_once(str) -> None: ...  # LogAPIUsageOnceFromPython
def _demangle(str) -> str: ...  # c10::demangle

# Defined in `valgrind.h` and `callgrind.h` respecitively.
def _valgrind_supported_platform() -> _bool: ...  # NVALGRIND
def _valgrind_toggle() -> None: ...  # CALLGRIND_TOGGLE_COLLECT

has_openmp: _bool
has_mkl: _bool
has_lapack: _bool
has_cuda: _bool
has_mkldnn: _bool
has_cudnn: _bool
_GLIBCXX_USE_CXX11_ABI: _bool
default_generator: Generator

# Defined in torch/csrc/autograd/init.cpp
def _set_grad_enabled(enabled: _bool) -> None: ...
def is_grad_enabled() -> _bool: ...
def set_autocast_enabled(enabled: _bool) -> None: ...
def is_autocast_enabled() -> _bool: ...
def clear_autocast_cache() -> None: ...
def autocast_increment_nesting() -> _int: ...
def autocast_decrement_nesting() -> _int: ...
def set_anomaly_enabled(enabled: _bool) -> None: ...
def is_anomaly_enabled() -> _bool: ...

# Defined in torch/csrc/jit/python/script_init.cpp
class LoggerBase(object):
    ...

class NoopLogger(LoggerBase):
    ...

class LockingLogger(LoggerBase):
    ...

class AggregationType(Enum):
    SUM = 0
    AVG = 1

class FileCheck(object):
    # TODO (add more FileCheck signature)
    def check_source_highlighted(self, highlight: str) -> 'FileCheck': ...
    def run(self, test_string: str) -> None: ...
    ...

# Defined in torch/csrc/jit/python/init.cpp
class PyTorchFileReader(object):
    @overload
    def __init__(self, name: str) -> None: ...
    @overload
    def __init__(self, buffer: BinaryIO) -> None: ...
    def get_record(self, name: str) -> bytes: ...
    ...

class PyTorchFileWriter(object):
    @overload
    def __init__(self, name: str) -> None: ...
    @overload
    def __init__(self, buffer: BinaryIO) -> None: ...
    def write_record(self, name: str, data: bytes, size: _int) -> None: ...
    def write_end_of_file(self) -> None: ...
    ...

def _jit_get_inline_everything_mode() -> _bool: ...
def _jit_set_inline_everything_mode(enabled: _bool) -> None: ...
def _jit_pass_dce(Graph) -> None: ...
def _jit_pass_lint(Graph) -> None: ...

# Defined in torch/csrc/jit/python/python_custome_class.cpp
def _get_custom_class_python_wrapper(name: str, attr: str) -> Any: ...

# Defined in torch/csrc/Generator.cpp
class Generator(object):
    device: _device
    def __init__(self, device: Union[_device, str, None] = None) -> None: ...
    def get_state(self) -> Tensor: ...
    def set_state(self, _new_state: Tensor) -> Generator: ...
    def manual_seed(self, seed: _int) -> Generator: ...
    def seed(self) -> _int: ...
    def initial_seed(self) -> _int: ...

# Defined in torch/csrc/utils/init.cpp
class BenchmarkConfig(object):
    num_calling_threads: _int
    num_worker_threads: _int
    num_warmup_iters: _int
    num_iters: _int
    profiler_output_path: str

class BenchmarkExecutionStats(object):
    latency_avg_ms: _float
    num_iters: _int

class ThroughputBenchmark(object):
    def __init__(self, module: Any) -> None: ...
    def add_input(self, *args: Any, **kwargs: Any) -> None: ...
    def run_once(self, *args: Any, **kwargs: Any) -> Any: ...
    def benchmark(self, config: BenchmarkConfig) -> BenchmarkExecutionStats: ...

# IDK if these are actually exposed here, hope they are
${namedtuple_defs}

# Defined in torch/csrc/generic/Storage.cpp
${legacy_storage_base_hints}

# TODO: where
${legacy_class_hints}

# Defined in torch/csrc/autograd/python_engine.cpp
class _ImperativeEngine:
    ...

# Defined in torch/csrc/autograd/python_variable.cpp
class _TensorBase(object):
    requires_grad: _bool
    shape: Size
    data: Tensor
    names: List[str]
    device: _device
    dtype: _dtype
    layout: _layout
    real: Tensor
    imag: Tensor
    T: Tensor
    ndim: _int
    _version: _int
    _base: Optional[Tensor]
    grad_fn: Any
    ${tensor_method_hints}

# Defined in torch/csrc/cuda/Module.cpp
def _cuda_getCurrentStream(device: _int) -> _int: ...
def _cuda_getDefaultStream(device: _int) -> _int: ...
def _cuda_getCurrentBlasHandle() -> _int: ...
def _cuda_setDevice(device: _int) -> None: ...
def _cuda_setStream(cuda_stream: _int) -> None: ...
def _cuda_getCompiledVersion() -> _int: ...
def _cuda_cudaHostAllocator() -> _int: ...
def _cuda_cudaCachingAllocator_raw_alloc(size: _int, cuda_stream: _int) -> _int: ...
def _cuda_cudaCachingAllocator_raw_delete(ptr: _int) -> None: ...
def _cuda_emptyCache() -> None: ...
def _cuda_memoryStats(device: _int) -> Dict[str, Any]: ...
def _cuda_resetAccumulatedMemoryStats(device: _int) -> None: ...
def _cuda_resetPeakMemoryStats(device: _int) -> None: ...
def _cuda_memorySnapshot() -> List[Dict[str, Any]]: ...
def _cuda_lock_mutex() -> None: ...
def _cuda_unlock_mutex() -> None: ...
def _nccl_version() -> _int: ...
def _nccl_unique_id() -> bytes: ...
def _nccl_init_rank(nranks: _int, comm_id: bytes, rank: _int) -> object: ...
def _nccl_reduce(input: Sequence[Tensor],
                 output: Tensor,
                 root: _int,
                 op: _int,
                 streams: Optional[Sequence[_CudaStreamBase]],
                 comms: Optional[Sequence[object]]) -> None: ...
def _nccl_all_reduce(input: Sequence[Tensor],
                     output: Sequence[Tensor],
                     op: _int,
                     streams: Optional[Sequence[_CudaStreamBase]],
                     comms: Optional[Sequence[object]]) -> None: ...
def _nccl_broadcast(input: Sequence[Tensor],
                    root: _int,
                    streams: Optional[Sequence[_CudaStreamBase]],
                    comms: Optional[Sequence[object]]) -> None: ...
def _nccl_all_gather(input: Sequence[Tensor],
                     output: Sequence[Tensor],
                     streams: Optional[Sequence[_CudaStreamBase]],
                     comms: Optional[Sequence[object]]) -> None: ...
def _nccl_reduce_scatter(input: Sequence[Tensor],
                         output: Sequence[Tensor],
                         op: _int,
                         streams: Optional[Sequence[_CudaStreamBase]],
                         comms: Optional[Sequence[object]]) -> None: ...


class _CudaDeviceProperties:
    name: str
    major: _int
    minor: _int
    multi_processor_count: _int
    total_memory: _int
    is_integrated: _int
    is_multi_gpu_board: _int

# Defined in torch/csrc/cuda/python_comm.cpp
def _broadcast(tensor: Tensor, devices: List[_int]) -> List[Tensor]: ...
def _broadcast_out(tensor: Tensor, out_tensors: List[Tensor]) -> List[Tensor]: ...
def _broadcast_coalesced(
    tensors: List[Tensor],
    devices: List[_int],
    buffer_size: _int
) -> List[List[Tensor]]: ...

def _scatter(tensor: Tensor, devices: List[_int], chunk_sizes: Optional[List[_int]], dim: _int, streams: Optional[List[Stream]]) -> List[Tensor]: ...
def _scatter_out(tensor: Tensor, out_tensors: List[Tensor], dim: _int, streams: Optional[List[Stream]]) -> List[Tensor]: ...
def _gather(tensors: List[Tensor], dim: _int, destination_index: Optional[_int]) -> Tensor: ...
def _gather_out(tensors: List[Tensor], out_tensor: Tensor, dim: _int) -> Tensor: ...

# Defined in torch/csrc/cuda/Stream.cpp
class _CudaStreamBase:
    _cdata: _int
    device: _device
    cuda_stream: _int
    priority: _int

    def __new__(self, priority: _int = 0, _cdata: _int = 0) -> _CudaStreamBase: ...
    def query(self) -> _bool: ...
    def synchronize(self) -> None: ...
    def priority_range(self) -> Tuple[_int, _int]: ...

# Defined in torch/csrc/cuda/Event.cpp
class _CudaEventBase:
    device: _device
    cuda_event: _int

    def __new__(cls, enable_timing: _bool = False, blocking: _bool = False, interprocess: _bool = False) -> _CudaEventBase: ...
    @classmethod
    def from_ipc_handle(cls, device: _device, ipc_handle: bytes) -> _CudaEventBase: ...
    def record(self, stream: _CudaStreamBase) -> None: ...
    def wait(self, stream: _CudaStreamBase) -> None: ...
    def query(self) -> _bool: ...
    def elapsed_time(self, other: _CudaEventBase) -> _float: ...
    def synchronize(self) -> None: ...
    def ipc_handle(self) -> bytes: ...

# Defined in torch/csrc/DataLoader.cpp
def _set_worker_signal_handlers(*arg: Any) -> None: ...  # THPModule_setWorkerSignalHandlers
def _set_worker_pids(key: _int, child_pids: Tuple[_int, ...]) -> None: ...  # THPModule_setWorkerPIDs
def _remove_worker_pids(loader_id: _int) -> None: ...  # THPModule_removeWorkerPIDs
def _error_if_any_worker_fails() -> None: ...  # THPModule_errorIfAnyWorkerFails

# Defined in torch/csrc/jit/python/python_tracer.cpp
class TracingState: ...

def _create_graph_by_tracing(
    func: Callable[..., Any],
    inputs: Any,
    var_name_lookup_fn: Callable[[Tensor], str],
    strict: Any,
    force_outplace: Any,
    self: Any = None
) -> Tuple[Graph, Stack]: ...
def _tracer_warn_use_python(): ...
def _get_tracing_state() -> TracingState: ...

# Defined in torch/csrc/jit/python/python_ir.cpp
# Not actually defined in python_ir.cpp, not sure where they are.
class IValue:
    ...
Stack = List[IValue]

class JitType:
    ...

R = TypeVar('R', bound=JitType)

class AnyType(JitType):
    @staticmethod
    def get() -> AnyType: ...

class NoneType(JitType):
    @staticmethod
    def get() -> NoneType: ...

class BoolType(JitType):
    @staticmethod
    def get() -> BoolType: ...

class FloatType(JitType):
    @staticmethod
    def get() -> FloatType: ...

class IntType(JitType):
    @staticmethod
    def get() -> IntType: ...

class StringType(JitType):
    @staticmethod
    def get() -> StringType: ...

class DeviceObjType(JitType):
    @staticmethod
    def get() -> DeviceObjType: ...

class StreamObjType(JitType):
    @staticmethod
    def get() -> StreamObjType: ...

class ListType(JitType):
    def __init__(self, a: JitType) -> None: ...
    def getElementType(self) -> JitType: ...

    @staticmethod
    def ofInts() -> ListType: ...
    @staticmethod
    def ofTensors() -> ListType: ...
    @staticmethod
    def ofFloats() -> ListType: ...
    @staticmethod
    def ofBools() -> ListType: ...

class DictType(JitType):
    def __init__(self, key: JitType, value: JitType) -> None: ...
    def getKeyType(self) -> JitType: ...
    def getValueType(self) -> JitType: ...

class TupleType(JitType):
    def __init__(self, a: List[JitType]) -> None: ...

class ClassType(JitType):
    def __init__(self, qualified_name: str) -> None: ...

class InterfaceType(JitType):
    def __init__(self, qualified_name: str) -> None: ...
    def getMethod(self, name: str) -> Optional[FunctionSchema]: ...
    def getMethodNames(self) -> List[str]: ...

class OptionalType(JitType, Generic[R]):
    def __init__(self, a: JitType) -> None: ...
    def getElementType(self) -> JitType: ...

    @staticmethod
    def ofTensor() -> OptionalType: ...

class FutureType(JitType):
    def __init__(self, a: JitType) -> None: ...
    def getElementType(self) -> JitType: ...

class RRefType(JitType):
    def __init__(self, a: JitType) -> None: ...

class EnumType(JitType):
    def __init__(
        self,
        qualified_name: str,
        value_type: JitType,
        enum_names_values: List[Any]
    ) -> None:
        ...

class TensorType(JitType):
    @classmethod
    def get(cls) -> TensorType: ...
    @classmethod
    def getInferred(cls) -> TensorType: ...

# Defined in torch/csrc/jit/python/python_tree_views.cpp
class SourceRange:
    ...

class TreeView:
    ...

class Ident(TreeView):
    @property
    def name(self) -> str: ...

class ClassDef(TreeView):
    ...

class Def(TreeView):
    def name(self) -> Ident: ...

class Decl(TreeView):
    ...

# Defined in torch/csrc/distributed/rpc/init.cpp
def _rpc_init() -> _bool: ...

# Defined in torch/csrc/distributed/autograd/init.cpp
def _dist_autograd_init() -> _bool: ...

# Defined in torch/csrc/distributed/c10d/init.cpp
def _c10d_init() -> _bool: ...

# Defined in torch/csrc/distributed/rpc/testing/init.cpp
def _faulty_agent_init() -> _bool: ...