Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,15 @@
|
||||
from torch import Tensor
|
||||
from torch.types import _dtype, _int, Device
|
||||
|
||||
# Defined in torch/csrc/acc/Module.cpp
|
||||
class PrivateUse1Hooks:
|
||||
def has_primary_context(self, device_index: _int) -> bool: ...
|
||||
def is_built(self) -> bool: ...
|
||||
def is_available(self) -> bool: ...
|
||||
|
||||
class DeviceGuard:
|
||||
def type_(self) -> Device: ...
|
||||
|
||||
def register_python_privateuseone_device_guard(guard: DeviceGuard) -> bool: ...
|
||||
def register_python_privateuseone_hook(hook: PrivateUse1Hooks) -> bool: ...
|
||||
def create_empty_tensor(shape: tuple[_int, ...], dtype: _dtype) -> Tensor: ...
|
||||
@@ -0,0 +1,164 @@
|
||||
from ctypes import c_void_p
|
||||
from typing import overload, Protocol
|
||||
|
||||
from torch import Tensor
|
||||
|
||||
# Defined in torch/csrc/inductor/aoti_runner/pybind.cpp
|
||||
|
||||
# Tensor to AtenTensorHandle
|
||||
def unsafe_alloc_void_ptrs_from_tensors(tensors: list[Tensor]) -> list[c_void_p]: ...
|
||||
def unsafe_alloc_void_ptr_from_tensor(tensor: Tensor) -> c_void_p: ...
|
||||
|
||||
# AtenTensorHandle to Tensor
|
||||
def alloc_tensors_by_stealing_from_void_ptrs(
|
||||
handles: list[c_void_p],
|
||||
) -> list[Tensor]: ...
|
||||
def alloc_tensor_by_stealing_from_void_ptr(
|
||||
handle: c_void_p,
|
||||
) -> Tensor: ...
|
||||
|
||||
class AOTIModelContainerRunner(Protocol):
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_names_to_original_fqns(self) -> dict[str, str]: ...
|
||||
def get_constant_names_to_dtypes(self) -> dict[str, int]: ...
|
||||
def extract_constants_map(self, use_inactive: bool) -> dict[str, Tensor]: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def swap_constant_buffer(self) -> None: ...
|
||||
def free_inactive_constant_buffer(self) -> None: ...
|
||||
|
||||
class AOTIModelContainerRunnerCpu:
|
||||
def __init__(self, model_so_path: str, num_models: int) -> None: ...
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_names_to_original_fqns(self) -> dict[str, str]: ...
|
||||
def get_constant_names_to_dtypes(self) -> dict[str, int]: ...
|
||||
def extract_constants_map(self, use_inactive: bool) -> dict[str, Tensor]: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def swap_constant_buffer(self) -> None: ...
|
||||
def free_inactive_constant_buffer(self) -> None: ...
|
||||
|
||||
class AOTIModelContainerRunnerCuda:
|
||||
@overload
|
||||
def __init__(self, model_so_path: str, num_models: int) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self, model_so_path: str, num_models: int, device_str: str
|
||||
) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self, model_so_path: str, num_models: int, device_str: str, cubin_dir: str
|
||||
) -> None: ...
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_names_to_original_fqns(self) -> dict[str, str]: ...
|
||||
def get_constant_names_to_dtypes(self) -> dict[str, int]: ...
|
||||
def extract_constants_map(self, use_inactive: bool) -> dict[str, Tensor]: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def swap_constant_buffer(self) -> None: ...
|
||||
def free_inactive_constant_buffer(self) -> None: ...
|
||||
|
||||
class AOTIModelContainerRunnerXpu:
|
||||
@overload
|
||||
def __init__(self, model_so_path: str, num_models: int) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self, model_so_path: str, num_models: int, device_str: str
|
||||
) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self, model_so_path: str, num_models: int, device_str: str, kernel_bin_dir: str
|
||||
) -> None: ...
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_names_to_original_fqns(self) -> dict[str, str]: ...
|
||||
def get_constant_names_to_dtypes(self) -> dict[str, int]: ...
|
||||
def extract_constants_map(self, use_inactive: bool) -> dict[str, Tensor]: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def swap_constant_buffer(self) -> None: ...
|
||||
def free_inactive_constant_buffer(self) -> None: ...
|
||||
|
||||
class AOTIModelContainerRunnerMps:
|
||||
def __init__(self, model_so_path: str, num_models: int) -> None: ...
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_names_to_original_fqns(self) -> dict[str, str]: ...
|
||||
def get_constant_names_to_dtypes(self) -> dict[str, int]: ...
|
||||
def extract_constants_map(self, use_inactive: bool) -> dict[str, Tensor]: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def swap_constant_buffer(self) -> None: ...
|
||||
def free_inactive_constant_buffer(self) -> None: ...
|
||||
|
||||
# Defined in torch/csrc/inductor/aoti_package/pybind.cpp
|
||||
class AOTIModelPackageLoader:
|
||||
def __init__(
|
||||
self,
|
||||
model_package_path: str,
|
||||
model_name: str,
|
||||
run_single_threaded: bool,
|
||||
num_runners: int,
|
||||
device_index: int,
|
||||
) -> None: ...
|
||||
def get_metadata(self) -> dict[str, str]: ...
|
||||
def run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def boxed_run(
|
||||
self, inputs: list[Tensor], stream_handle: c_void_p = ...
|
||||
) -> list[Tensor]: ...
|
||||
def get_call_spec(self) -> list[str]: ...
|
||||
def get_constant_fqns(self) -> list[str]: ...
|
||||
def load_constants(
|
||||
self,
|
||||
constants_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
check_full_update: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
def update_constant_buffer(
|
||||
self,
|
||||
tensor_map: dict[str, Tensor],
|
||||
use_inactive: bool,
|
||||
validate_full_updates: bool,
|
||||
user_managed: bool = ...,
|
||||
) -> None: ...
|
||||
@@ -0,0 +1,163 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Callable
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch._C._profiler import (
|
||||
_ProfilerEvent,
|
||||
ActiveProfilerType,
|
||||
ProfilerActivity,
|
||||
ProfilerConfig,
|
||||
)
|
||||
|
||||
# Defined in torch/csrc/autograd/init.cpp
|
||||
|
||||
class DeviceType(Enum):
|
||||
CPU = ...
|
||||
CUDA = ...
|
||||
XPU = ...
|
||||
MKLDNN = ...
|
||||
OPENGL = ...
|
||||
OPENCL = ...
|
||||
IDEEP = ...
|
||||
HIP = ...
|
||||
FPGA = ...
|
||||
MAIA = ...
|
||||
XLA = ...
|
||||
MTIA = ...
|
||||
MPS = ...
|
||||
HPU = ...
|
||||
Meta = ...
|
||||
Vulkan = ...
|
||||
Metal = ...
|
||||
PrivateUse1 = ...
|
||||
|
||||
class ProfilerEvent:
|
||||
def cpu_elapsed_us(self, other: ProfilerEvent) -> float: ...
|
||||
def cpu_memory_usage(self) -> int: ...
|
||||
def cuda_elapsed_us(self, other: ProfilerEvent) -> float: ...
|
||||
def privateuse1_elapsed_us(self, other: ProfilerEvent) -> float: ...
|
||||
def cuda_memory_usage(self) -> int: ...
|
||||
def device(self) -> int: ...
|
||||
def handle(self) -> int: ...
|
||||
def has_cuda(self) -> bool: ...
|
||||
def is_remote(self) -> bool: ...
|
||||
def kind(self) -> int: ...
|
||||
def name(self) -> str: ...
|
||||
def node_id(self) -> int: ...
|
||||
def sequence_nr(self) -> int: ...
|
||||
def shapes(self) -> list[list[int]]: ...
|
||||
def thread_id(self) -> int: ...
|
||||
def flops(self) -> float: ...
|
||||
def is_async(self) -> bool: ...
|
||||
|
||||
class _KinetoEvent:
|
||||
def name(self) -> str: ...
|
||||
def overload_name(self) -> str: ...
|
||||
def device_index(self) -> int: ...
|
||||
def device_resource_id(self) -> int: ...
|
||||
def start_ns(self) -> int: ...
|
||||
def end_ns(self) -> int: ...
|
||||
def duration_ns(self) -> int: ...
|
||||
def is_async(self) -> bool: ...
|
||||
def linked_correlation_id(self) -> int: ...
|
||||
def external_id(self) -> int: ...
|
||||
def shapes(self) -> list[list[int]]: ...
|
||||
def dtypes(self) -> list[str]: ...
|
||||
def concrete_inputs(self) -> list[Any]: ...
|
||||
def kwinputs(self) -> dict[str, Any]: ...
|
||||
def device_type(self) -> DeviceType: ...
|
||||
def start_thread_id(self) -> int: ...
|
||||
def end_thread_id(self) -> int: ...
|
||||
def correlation_id(self) -> int: ...
|
||||
def fwd_thread_id(self) -> int: ...
|
||||
def stack(self) -> list[str]: ...
|
||||
def scope(self) -> int: ...
|
||||
def sequence_nr(self) -> int: ...
|
||||
def flops(self) -> int: ...
|
||||
def cuda_elapsed_us(self) -> int: ...
|
||||
def privateuse1_elapsed_us(self) -> int: ...
|
||||
def is_user_annotation(self) -> bool: ...
|
||||
def is_python_function(self) -> bool: ...
|
||||
def is_hidden_event(self) -> bool: ...
|
||||
def metadata_json(self) -> str: ...
|
||||
def activity_type(self) -> str: ...
|
||||
def extra_meta(self) -> dict[str, str]: ...
|
||||
def flow_id(self) -> int: ...
|
||||
def flow_type(self) -> int: ...
|
||||
def flow_start(self) -> bool: ...
|
||||
def structured_input_shapes(self) -> list[list[int] | list[list[int]]]: ...
|
||||
def structured_input_strides(self) -> list[list[int] | list[list[int]]]: ...
|
||||
def python_id(self) -> int: ...
|
||||
def python_parent_id(self) -> int: ...
|
||||
def python_module_id(self) -> int: ...
|
||||
|
||||
class _ProfilerResult:
|
||||
def events(self) -> list[_KinetoEvent]: ...
|
||||
def legacy_events(self) -> list[list[ProfilerEvent]]: ...
|
||||
def save(self, path: str) -> None: ...
|
||||
def experimental_event_tree(self) -> list[_ProfilerEvent]: ...
|
||||
def trace_start_ns(self) -> int: ...
|
||||
|
||||
class SavedTensor:
|
||||
def unpack(self) -> torch.Tensor: ...
|
||||
|
||||
def _make_saved_tensor(
|
||||
tensor: torch.Tensor,
|
||||
is_output: bool,
|
||||
is_inplace_on_view: bool = False,
|
||||
) -> SavedTensor: ...
|
||||
def _enable_profiler(
|
||||
config: ProfilerConfig,
|
||||
activities: set[ProfilerActivity],
|
||||
) -> None: ...
|
||||
def _prepare_profiler(
|
||||
config: ProfilerConfig,
|
||||
activities: set[ProfilerActivity],
|
||||
activity_filter: dict[ProfilerActivity, set[str]] = ...,
|
||||
) -> None: ...
|
||||
def _toggle_collection_dynamic(
|
||||
enable: bool,
|
||||
activities: set[ProfilerActivity],
|
||||
) -> None: ...
|
||||
def _disable_profiler() -> _ProfilerResult: ...
|
||||
def _profiler_enabled() -> bool: ...
|
||||
def _add_metadata_json(key: str, value: str) -> None: ...
|
||||
def _kineto_step() -> None: ...
|
||||
def _get_current_graph_task_keep_graph() -> bool: ...
|
||||
def _get_sequence_nr() -> int: ...
|
||||
def kineto_available() -> bool: ...
|
||||
def _record_function_with_args_enter(name: str, *args) -> torch.Tensor: ...
|
||||
def _record_function_with_args_exit(handle: torch.Tensor) -> None: ...
|
||||
def _supported_activities() -> set[ProfilerActivity]: ...
|
||||
def _enable_record_function(enable: bool) -> None: ...
|
||||
def _set_empty_test_observer(is_global: bool, sampling_prob: float) -> None: ...
|
||||
def _push_saved_tensors_default_hooks(
|
||||
pack_hook: Callable[[torch.Tensor], Any],
|
||||
unpack_hook: Callable[[Any], torch.Tensor],
|
||||
) -> None: ...
|
||||
def _pop_saved_tensors_default_hooks() -> None: ...
|
||||
def _top_saved_tensors_default_hooks(
|
||||
ignore_is_tracing: bool,
|
||||
) -> tuple[Callable[[torch.Tensor], Any], Callable[[Any], torch.Tensor]]: ...
|
||||
def _unsafe_set_version_counter(
|
||||
t: tuple[torch.Tensor, ...], prev_version: tuple[int, ...]
|
||||
) -> None: ...
|
||||
def _enable_profiler_legacy(config: ProfilerConfig) -> None: ...
|
||||
def _disable_profiler_legacy() -> list[list[ProfilerEvent]]: ...
|
||||
def _profiler_type() -> ActiveProfilerType: ...
|
||||
def _saved_tensors_hooks_enable() -> None: ...
|
||||
def _saved_tensors_hooks_disable(message: str, fail_if_non_empty=True) -> None: ...
|
||||
def _saved_tensors_hooks_get_disabled_error_message() -> str | None: ...
|
||||
def _saved_tensors_hooks_set_tracing(is_tracing: bool) -> bool: ...
|
||||
|
||||
class CreationMeta(Enum):
|
||||
DEFAULT = ...
|
||||
IN_CUSTOM_FUNCTION = ...
|
||||
MULTI_OUTPUT_NODE = ...
|
||||
NO_GRAD_MODE = ...
|
||||
INFERENCE_MODE = ...
|
||||
|
||||
def _set_creation_meta(t: torch.Tensor, creation_meta: CreationMeta) -> None: ...
|
||||
def _get_creation_meta(t: torch.Tensor) -> CreationMeta: ...
|
||||
@@ -0,0 +1,8 @@
|
||||
from typing import Any
|
||||
|
||||
from torch.types import _bool
|
||||
|
||||
# Defined in torch/csrc/cpu/Module.cpp
|
||||
|
||||
def _init_amx() -> _bool: ...
|
||||
def _get_cpu_capability() -> dict[str, Any]: ...
|
||||
@@ -0,0 +1,14 @@
|
||||
from enum import IntEnum
|
||||
|
||||
# Defined in torch/csrc/cuda/shared/cudnn.cpp
|
||||
is_cuda: bool
|
||||
|
||||
def getRuntimeVersion() -> tuple[int, int, int]: ...
|
||||
def getCompileVersion() -> tuple[int, int, int]: ...
|
||||
def getVersionInt() -> int: ...
|
||||
|
||||
class RNNMode(IntEnum):
|
||||
rnn_relu = ...
|
||||
rnn_tanh = ...
|
||||
lstm = ...
|
||||
gru = ...
|
||||
@@ -0,0 +1 @@
|
||||
def getVersionInt() -> int: ...
|
||||
@@ -0,0 +1,21 @@
|
||||
# This module is defined in torch/csrc/distributed/python_placement.cpp
|
||||
|
||||
class Placement:
|
||||
def is_partial(self, reduce_op: str | None = None) -> bool: ...
|
||||
def is_replicate(self) -> bool: ...
|
||||
def is_shard(self, dim: int | None = None) -> bool: ...
|
||||
|
||||
class Shard(Placement):
|
||||
dim: int
|
||||
def __init__(self, dim: int): ...
|
||||
|
||||
class StridedShard(Placement):
|
||||
dim: int
|
||||
split_factor: int
|
||||
def __init__(self, dim: int, *, split_factor: int): ...
|
||||
|
||||
class Replicate(Placement): ...
|
||||
|
||||
class Partial(Placement):
|
||||
reduce_op: str
|
||||
def __init__(self, reduce_op: str | None = None): ...
|
||||
@@ -0,0 +1,26 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
# This module is defined in torch/csrc/distributed/autograd/init.cpp
|
||||
|
||||
class DistAutogradContext:
|
||||
def _context_id(self) -> int: ...
|
||||
def _recv_functions(self) -> dict[int, Any]: ...
|
||||
def _send_functions(self) -> dict[int, Any]: ...
|
||||
def _known_worker_ids(self) -> set[int]: ...
|
||||
|
||||
def _new_context() -> DistAutogradContext: ...
|
||||
def _release_context(context_id: int) -> None: ...
|
||||
def _get_max_id() -> int: ...
|
||||
def _is_valid_context(worker_id: int) -> bool: ...
|
||||
def _retrieve_context(context_id: int) -> DistAutogradContext: ...
|
||||
def _current_context() -> DistAutogradContext: ...
|
||||
def _init(worker_id: int) -> None: ...
|
||||
def _get_debug_info() -> dict[str, str]: ...
|
||||
def backward(
|
||||
context_id: int,
|
||||
roots: list[torch.Tensor],
|
||||
retain_graph: bool = False,
|
||||
) -> None: ...
|
||||
def get_gradients(context_id: int) -> dict[torch.Tensor, torch.Tensor]: ...
|
||||
@@ -0,0 +1,900 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# mypy: disable-error-code="type-arg"
|
||||
from collections.abc import Callable
|
||||
from datetime import timedelta
|
||||
from enum import Enum
|
||||
from typing import Any, overload
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch._C import ScriptObject
|
||||
from torch._C._autograd import DeviceType
|
||||
from torch.distributed.distributed_c10d import GroupName
|
||||
from torch.futures import Future
|
||||
|
||||
# This module is defined in torch/csrc/distributed/c10d/init.cpp
|
||||
|
||||
_DEFAULT_FIRST_BUCKET_BYTES: int
|
||||
_DEFAULT_NO_TIMEOUT: timedelta
|
||||
_DEFAULT_PG_TIMEOUT: timedelta
|
||||
_DEFAULT_PG_NCCL_TIMEOUT: timedelta
|
||||
|
||||
class BuiltinCommHookType(Enum):
|
||||
ALLREDUCE = ...
|
||||
FP16_COMPRESS = ...
|
||||
|
||||
def _register_comm_hook(reducer: Reducer, state: Any, comm_hook: Any): ...
|
||||
def _register_builtin_comm_hook(
|
||||
reducer: Reducer,
|
||||
comm_hook_type: BuiltinCommHookType,
|
||||
): ...
|
||||
def _set_global_rank(rank: int) -> None: ...
|
||||
def _hash_tensors(tensors: list[Tensor]) -> int: ...
|
||||
|
||||
class GradBucket:
|
||||
def index(self) -> int: ...
|
||||
def buffer(self) -> Tensor: ...
|
||||
def gradients(self) -> list[Tensor]: ...
|
||||
def is_last(self) -> bool: ...
|
||||
def set_buffer(self, tensor: Tensor) -> None: ...
|
||||
def parameters(self) -> list[Tensor]: ...
|
||||
|
||||
class Reducer:
|
||||
def __init__(
|
||||
self,
|
||||
params: list[Tensor],
|
||||
bucket_indices: list[list[int]],
|
||||
per_bucket_size_limits: list[int],
|
||||
process_group: ProcessGroup,
|
||||
expect_sparse_gradients: list[bool] = ...,
|
||||
bucket_bytes_cap: int = ..., # kDefaultBucketBytesCap in reducer.hpp
|
||||
find_unused_parameters: bool = ...,
|
||||
gradient_as_bucket_view: bool = ...,
|
||||
param_to_name_mapping: dict[int, str] = ...,
|
||||
first_bucket_types_cap: int = ..., # kDefaultFirstBucketBytes in reducer.hpp
|
||||
skip_all_reduce_unused_params: bool = ...,
|
||||
use_python_reducer: bool = ...,
|
||||
bucket_bytes_cap_list: list[int] = ...,
|
||||
batched_grad_copy: bool = ...,
|
||||
) -> None: ...
|
||||
def prepare_for_forward(self) -> None: ...
|
||||
def prepare_for_backward(self, output: list[Tensor]) -> None: ...
|
||||
def get_backward_stats(self) -> list[int]: ...
|
||||
def _install_post_backward_futures(self, futures: list[Future]) -> None: ...
|
||||
def _rebuild_buckets(self) -> bool: ...
|
||||
def _get_zeros_like_grad_buckets(self) -> list[GradBucket]: ...
|
||||
def _push_all_rebuilt_params(self) -> None: ...
|
||||
def _set_forward_pass_work_handle(
|
||||
self,
|
||||
work: Work,
|
||||
use_static_world_size: bool,
|
||||
): ...
|
||||
def _get_local_used_map(self) -> Tensor: ...
|
||||
def _set_ddp_runtime_logging_sample_rate(self, sample_rate: int) -> None: ...
|
||||
def _set_static_graph(self) -> None: ...
|
||||
def _run_comm_hook(self, bucket: GradBucket) -> Future: ...
|
||||
def set_logger(self, logger: Logger) -> None: ...
|
||||
def _remove_autograd_hooks(self) -> None: ...
|
||||
def _check_reducer_finalized(self) -> None: ...
|
||||
def _set_sparse_metadata(self, global_unique_ids: dict[str, Tensor]) -> None: ...
|
||||
def _reset_state(self) -> None: ...
|
||||
def _update_process_group(self, new_process_group: ProcessGroup) -> None: ...
|
||||
|
||||
class DDPLoggingData:
|
||||
strs_map: dict[str, str]
|
||||
ints_map: dict[str, int]
|
||||
|
||||
class Logger:
|
||||
def __init__(self, reducer: Reducer) -> None: ...
|
||||
def set_construction_data_and_log(
|
||||
self,
|
||||
module_name: str,
|
||||
device_ids: list[int],
|
||||
output_device: int,
|
||||
broadcast_buffers: bool,
|
||||
has_sync_bn: bool,
|
||||
static_graph: bool,
|
||||
) -> None: ...
|
||||
def set_runtime_stats_and_log(self) -> None: ...
|
||||
def set_error_and_log(self, error: str) -> None: ...
|
||||
def _get_ddp_logging_data(self) -> DDPLoggingData: ...
|
||||
def _set_comm_hook_name(self, comm_hook: str) -> None: ...
|
||||
def _set_uneven_input_join(self) -> None: ...
|
||||
def _set_static_graph(self) -> None: ...
|
||||
|
||||
class _WorkerServer:
|
||||
port: int
|
||||
|
||||
def __init__(self, host_or_file: str, port: int = ...) -> None: ...
|
||||
def shutdown(self) -> None: ...
|
||||
|
||||
class DebugLevel(Enum):
|
||||
OFF = ...
|
||||
INFO = ...
|
||||
DETAIL = ...
|
||||
|
||||
def get_debug_level() -> DebugLevel: ...
|
||||
def set_debug_level(level: DebugLevel) -> None: ...
|
||||
def set_debug_level_from_env() -> None: ...
|
||||
|
||||
class ReduceOp:
|
||||
# pyrefly: ignore # unknown-name
|
||||
def __init__(self, op: RedOpType) -> None: ...
|
||||
|
||||
# pyrefly: ignore # unknown-name
|
||||
SUM: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
AVG: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
PRODUCT: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
MIN: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
MAX: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
BAND: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
BOR: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
BXOR: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
PREMUL_SUM: RedOpType = ...
|
||||
# pyrefly: ignore # unknown-name
|
||||
UNUSED: RedOpType = ...
|
||||
|
||||
# mypy error being ignored:
|
||||
# Detected enum "torch._C._distributed_c10d.ReduceOp.RedOpType" in a type
|
||||
# stub with zero members. There is a chance this is due to a recent change
|
||||
# in the semantics of enum membership. If so, use `member = value` to mark
|
||||
# an enum member, instead of `member: type`
|
||||
class RedOpType(Enum): ... # type: ignore[misc]
|
||||
|
||||
class BroadcastOptions:
|
||||
rootRank: int
|
||||
rootTensor: int
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class AllreduceOptions:
|
||||
reduceOp: ReduceOp
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
sparseIndices: Tensor | None
|
||||
|
||||
class AllreduceCoalescedOptions(AllreduceOptions): ...
|
||||
|
||||
class ReduceOptions:
|
||||
reduceOp: ReduceOp
|
||||
rootRank: int
|
||||
rootTensor: int
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class AllgatherOptions:
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class GatherOptions:
|
||||
rootRank: int
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class ScatterOptions:
|
||||
rootRank: int
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class ReduceScatterOptions:
|
||||
reduceOp: ReduceOp
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class BarrierOptions:
|
||||
device_ids: list[int]
|
||||
device: torch.device
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class AllToAllOptions:
|
||||
timeout: timedelta
|
||||
asyncOp: bool
|
||||
|
||||
class Store:
|
||||
def set(self, key: str, value: str) -> None: ...
|
||||
def get(self, key: str) -> bytes: ...
|
||||
def add(self, key: str, value: int) -> int: ...
|
||||
def check(self, keys: list[str]) -> bool: ...
|
||||
def compare_set(
|
||||
self,
|
||||
key: str,
|
||||
expected_value: str,
|
||||
desired_value: str,
|
||||
) -> bytes: ...
|
||||
def delete_key(self, key: str) -> bool: ...
|
||||
def multi_get(self, keys: list[str]) -> list[bytes]: ...
|
||||
def num_keys(self) -> int: ...
|
||||
def set_timeout(self, timeout: timedelta) -> None: ...
|
||||
@overload
|
||||
def wait(self, keys: list[str]) -> None: ...
|
||||
@overload
|
||||
def wait(self, keys: list[str], timeout: timedelta) -> None: ...
|
||||
def queue_pop(self, key: str, block: bool = True) -> bytes: ...
|
||||
def queue_push(self, key: str, value: bytes | str) -> None: ...
|
||||
def queue_len(self, key: str) -> int: ...
|
||||
def list_keys(self) -> list[str]: ...
|
||||
|
||||
class FileStore(Store):
|
||||
def __init__(self, path: str, numWorkers: int = ...) -> None: ...
|
||||
|
||||
class HashStore(Store):
|
||||
def __init__(self) -> None: ...
|
||||
|
||||
class TCPStore(Store):
|
||||
def __init__(
|
||||
self,
|
||||
host_name: str,
|
||||
port: int,
|
||||
world_size: int | None = ...,
|
||||
is_master: bool = ...,
|
||||
timeout: timedelta = ...,
|
||||
wait_for_workers: bool = ...,
|
||||
multi_tenant: bool = ...,
|
||||
master_listen_fd: int | None = ...,
|
||||
use_libuv: bool | None = ...,
|
||||
) -> None: ...
|
||||
@property
|
||||
def host(self) -> str: ...
|
||||
@property
|
||||
def port(self) -> int: ...
|
||||
|
||||
class PrefixStore(Store):
|
||||
def __init__(self, prefix: str, store: Store) -> None: ...
|
||||
@property
|
||||
def underlying_store(self) -> Store: ...
|
||||
|
||||
class _ControlCollectives:
|
||||
def barrier(self, key: str, timeout: timedelta, blocking: bool) -> None: ...
|
||||
def broadcast_send(self, key: str, data: str, timeout: timedelta) -> None: ...
|
||||
def broadcast_recv(self, key: str, timeout: timedelta) -> str: ...
|
||||
def gather_send(self, key: str, data: str, timeout: timedelta) -> None: ...
|
||||
def gather_recv(self, key: str, timeout: timedelta) -> str: ...
|
||||
def scatter_send(self, key: str, data: str, timeout: timedelta) -> None: ...
|
||||
def scatter_recv(self, key: str, timeout: timedelta) -> str: ...
|
||||
def all_gather(self, key: str, data: str, timeout: timedelta) -> str: ...
|
||||
def all_sum(self, key: str, data: int, timeout: timedelta) -> int: ...
|
||||
|
||||
class _StoreCollectives(_ControlCollectives):
|
||||
def __init__(self, store: Store, rank: int, world_size: int) -> None: ...
|
||||
|
||||
class _DistributedBackendOptions:
|
||||
def __init__(self) -> None: ...
|
||||
@property
|
||||
def store(self) -> Store: ...
|
||||
@store.setter
|
||||
def store(self, store: Store) -> None: ...
|
||||
@property
|
||||
def group_rank(self) -> int: ...
|
||||
@group_rank.setter
|
||||
def group_rank(self, rank: int) -> None: ...
|
||||
@property
|
||||
def group_size(self) -> int: ...
|
||||
@group_size.setter
|
||||
def group_size(self, size: int) -> None: ...
|
||||
@property
|
||||
def timeout(self) -> timedelta: ...
|
||||
@timeout.setter
|
||||
def timeout(self, timeout: timedelta) -> None: ...
|
||||
@property
|
||||
def group_id(self) -> str: ...
|
||||
@group_id.setter
|
||||
def group_id(self, group_id: str) -> None: ...
|
||||
@property
|
||||
def global_ranks_in_group(self) -> list[int]: ...
|
||||
@global_ranks_in_group.setter
|
||||
def global_ranks_in_group(self, ranks: list[int]) -> None: ...
|
||||
|
||||
class Work:
|
||||
def is_completed(self) -> bool: ...
|
||||
def is_success(self) -> bool: ...
|
||||
def exception(self) -> Any: ...
|
||||
def wait(self, timeout: timedelta = ...) -> bool: ...
|
||||
def block_current_stream(self) -> None: ...
|
||||
def get_future(self) -> Future: ...
|
||||
def source_rank(self) -> int: ...
|
||||
def _source_rank(self) -> int: ...
|
||||
def result(self) -> list[Tensor]: ...
|
||||
def synchronize(self) -> None: ...
|
||||
def boxed(self) -> ScriptObject: ...
|
||||
@staticmethod
|
||||
def unbox(obj: ScriptObject) -> Work: ...
|
||||
|
||||
class Backend:
|
||||
class Options:
|
||||
def __init__(self, backend: str, timeout: timedelta = ...) -> None: ...
|
||||
@property
|
||||
def backend(self) -> str: ...
|
||||
@property
|
||||
def _timeout(self) -> timedelta: ...
|
||||
@_timeout.setter
|
||||
def _timeout(self, val: timedelta) -> None: ...
|
||||
global_ranks_in_group: list[int]
|
||||
group_name: GroupName
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rank: int,
|
||||
size: int,
|
||||
) -> None: ...
|
||||
@property
|
||||
def supports_splitting(self) -> bool: ...
|
||||
@property
|
||||
def supports_coalescing(self) -> bool: ...
|
||||
@property
|
||||
def supports_time_estimate(self) -> bool: ...
|
||||
def set_timeout(self, timeout: timedelta) -> None: ...
|
||||
@property
|
||||
def options(self) -> Options: ...
|
||||
def rank(self) -> int: ...
|
||||
def size(self) -> int: ...
|
||||
def name(self) -> str: ...
|
||||
def abort(self) -> None: ...
|
||||
def shutdown(self) -> None: ...
|
||||
def eager_connect_single_device(self, device: torch.device | None) -> None: ...
|
||||
def _set_sequence_number_for_group(self) -> None: ...
|
||||
def _set_default_timeout(self, timeout: timedelta) -> None: ...
|
||||
def get_error(self) -> ErrorType: ...
|
||||
def supports_tensor_alloc(self, device: torch.device) -> bool: ...
|
||||
def allocate_tensor(
|
||||
self,
|
||||
size: int,
|
||||
*,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> Tensor: ...
|
||||
@property
|
||||
def mem_allocator(self) -> Any: ...
|
||||
|
||||
class ProcessGroup:
|
||||
class BackendType(Enum):
|
||||
UNDEFINED = ...
|
||||
GLOO = ...
|
||||
NCCL = ...
|
||||
UCC = ...
|
||||
MPI = ...
|
||||
XCCL = ...
|
||||
CUSTOM = ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
rank: int,
|
||||
size: int,
|
||||
) -> None: ...
|
||||
def rank(self) -> int: ...
|
||||
def size(self) -> int: ...
|
||||
def get_group_store(self) -> Store: ...
|
||||
def split_group(
|
||||
self,
|
||||
new_ranks: list[int],
|
||||
timeout: timedelta | None = None,
|
||||
opts: Backend.Options | None = None,
|
||||
group_name: GroupName | None = None,
|
||||
group_desc: str | None = None,
|
||||
) -> ProcessGroup | None: ...
|
||||
def merge_remote_group(
|
||||
self,
|
||||
store: Store,
|
||||
size: int,
|
||||
timeout: timedelta,
|
||||
group_name: GroupName | None = None,
|
||||
group_desc: str | None = None,
|
||||
) -> ProcessGroup: ...
|
||||
def abort(self) -> None: ...
|
||||
def set_timeout(self, timeout: timedelta) -> None: ...
|
||||
def shutdown(self) -> None: ...
|
||||
@overload
|
||||
def broadcast(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def broadcast(
|
||||
self,
|
||||
tensor: Tensor,
|
||||
root: int,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def allreduce(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
opts: AllreduceOptions = ...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def allreduce(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
op=...,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def allreduce(
|
||||
self,
|
||||
tensor: Tensor,
|
||||
op=...,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
def allreduce_coalesced(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
def reduce_scatter_tensor_coalesced(
|
||||
self,
|
||||
outputTensors: list[Tensor],
|
||||
inputTensors: list[Tensor],
|
||||
opts: ReduceScatterOptions | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def reduce(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def reduce(
|
||||
self,
|
||||
tensor: Tensor,
|
||||
root: int,
|
||||
op=...,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def allgather(
|
||||
self,
|
||||
output_tensors: list[list[Tensor]],
|
||||
input_tensors: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def allgather(
|
||||
self,
|
||||
output_tensors: list[Tensor],
|
||||
input_tensor: Tensor,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
def _allgather_base(
|
||||
self,
|
||||
output: Tensor,
|
||||
input: Tensor,
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
def allgather_coalesced(
|
||||
self,
|
||||
output_lists: list[list[Tensor]],
|
||||
input_list: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
def allgather_into_tensor_coalesced(
|
||||
self,
|
||||
output_lists: list[Tensor],
|
||||
input_list: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def gather(
|
||||
self,
|
||||
output_tensors: list[list[Tensor]],
|
||||
input_tensors: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def gather(
|
||||
self,
|
||||
output_tensors: list[Tensor],
|
||||
input_tensor: Tensor,
|
||||
root: int,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def scatter(
|
||||
self,
|
||||
output_tensors: list[Tensor],
|
||||
input_tensors: list[list[Tensor]],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def scatter(
|
||||
self,
|
||||
output_tensor: Tensor,
|
||||
input_tensors: list[Tensor],
|
||||
root: int,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def reduce_scatter(
|
||||
self,
|
||||
output_tensors: list[Tensor],
|
||||
input_tensors: list[list[Tensor]],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def reduce_scatter(
|
||||
self,
|
||||
output_tensors: Tensor,
|
||||
input_tensor: list[Tensor],
|
||||
op=...,
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
def _reduce_scatter_base(
|
||||
self,
|
||||
outputTensor: Tensor,
|
||||
inputTensor: Tensor,
|
||||
opts: ReduceScatterOptions | None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def alltoall_base(
|
||||
self,
|
||||
output_tensor: Tensor,
|
||||
input_tensor: Tensor,
|
||||
output_split_sizes: list[int],
|
||||
input_split_sizes: list[int],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def alltoall_base(
|
||||
self,
|
||||
output: Tensor,
|
||||
input: Tensor,
|
||||
output_split_sizes: list[int],
|
||||
input_split_sizes: list[int],
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def alltoall(
|
||||
self,
|
||||
output_tensor: list[Tensor],
|
||||
input_tensor: list[Tensor],
|
||||
opts=...,
|
||||
) -> Work: ...
|
||||
@overload
|
||||
def alltoall(
|
||||
self,
|
||||
output: list[Tensor],
|
||||
input: list[Tensor],
|
||||
timeout: timedelta | None = None,
|
||||
) -> Work: ...
|
||||
def send(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
dstRank: int,
|
||||
tag: int,
|
||||
) -> Work: ...
|
||||
def recv(
|
||||
self,
|
||||
tensors: list[Tensor],
|
||||
srcRank: int,
|
||||
tag: int,
|
||||
) -> Work: ...
|
||||
def recv_anysource(self, tensors: list[Tensor], tag: int) -> Work: ...
|
||||
@overload
|
||||
def barrier(self, opts=...) -> Work: ...
|
||||
@overload
|
||||
def barrier(self, timeout: timedelta | None = None) -> Work: ...
|
||||
def boxed(self) -> ScriptObject: ...
|
||||
@staticmethod
|
||||
def unbox(obj: ScriptObject) -> ProcessGroup: ...
|
||||
def _start_coalescing(self, device: torch.device) -> None: ...
|
||||
def _end_coalescing(self, device: torch.device) -> Work: ...
|
||||
def _get_backend_name(self) -> str: ...
|
||||
def _backend_id(self, backend_type: BackendType) -> int: ...
|
||||
@property
|
||||
def _device_types(self) -> list[torch.device]: ...
|
||||
def _get_backend(self, device: torch.device) -> Backend: ...
|
||||
def _set_default_backend(self, backend_type: BackendType) -> None: ...
|
||||
def _register_backend(
|
||||
self,
|
||||
device: torch.device,
|
||||
backend_type: BackendType,
|
||||
backend: Backend | None,
|
||||
) -> None: ...
|
||||
def _set_group_name(self, name: GroupName) -> None: ...
|
||||
def _set_group_desc(self, desc: str) -> None: ...
|
||||
def name(self) -> str: ...
|
||||
def _has_hooks(self) -> bool: ...
|
||||
def _wait_for_pending_works(self) -> None: ...
|
||||
def _set_sequence_number_for_group(self) -> None: ...
|
||||
@property
|
||||
def bound_device_id(self) -> torch.device | None: ...
|
||||
@bound_device_id.setter
|
||||
def bound_device_id(self, device: torch.device | None) -> None: ...
|
||||
@property
|
||||
def group_name(self) -> GroupName: ...
|
||||
@property
|
||||
def group_desc(self) -> str: ...
|
||||
|
||||
class FakeProcessGroup(Backend):
|
||||
@staticmethod
|
||||
def _create_internal(rank: int, world_size: int) -> FakeProcessGroup: ...
|
||||
|
||||
class FakeWork(Work):
|
||||
seq_id: int
|
||||
def __init__(self) -> None: ...
|
||||
def wait(self, timeout: timedelta = ...) -> bool: ...
|
||||
def getFuture(self) -> Future: ...
|
||||
|
||||
class PythonCallbackWork(Work):
|
||||
def __init__(self, callback: Callable[[timedelta], bool]) -> None: ...
|
||||
def wait(self, timeout: timedelta = ...) -> bool: ...
|
||||
def get_future(self) -> Future: ...
|
||||
|
||||
class ProcessGroupGloo(Backend):
|
||||
class Device: ...
|
||||
|
||||
class Options(Backend.Options):
|
||||
devices: list[ProcessGroupGloo.Device]
|
||||
threads: int
|
||||
|
||||
def __init__(self): ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
rank: int,
|
||||
size: int,
|
||||
timeout: timedelta,
|
||||
) -> None: ...
|
||||
@staticmethod
|
||||
def create_device(hostname="", interface="", lazy_init=None) -> Device: ...
|
||||
@staticmethod
|
||||
def create_default_device(lazy_init=None) -> Device: ...
|
||||
def _set_default_timeout(self, timeout) -> None: ...
|
||||
@property
|
||||
def options(self) -> Options: ... # type: ignore[override]
|
||||
|
||||
class _ProcessGroupWrapper(Backend):
|
||||
def __init__(self, pg: Backend, gloo_pg: ProcessGroupGloo) -> None: ...
|
||||
wrapped_pg: Backend
|
||||
@property
|
||||
def options(self) -> Backend.Options: ...
|
||||
def get_error(self) -> ErrorType: ...
|
||||
|
||||
class ErrorType(Enum):
|
||||
SUCCESS = ...
|
||||
TIMEOUT = ...
|
||||
COMM_ERROR = ...
|
||||
REMOTE_ERROR = ...
|
||||
|
||||
class ProcessGroupNCCL(Backend):
|
||||
class NCCLConfig:
|
||||
blocking: int
|
||||
cga_cluster_size: int
|
||||
min_ctas: int
|
||||
max_ctas: int
|
||||
def unsafe_get_ptr(self) -> int: ...
|
||||
|
||||
class Options(Backend.Options):
|
||||
config: ProcessGroupNCCL.NCCLConfig
|
||||
is_high_priority_stream: bool
|
||||
split_from: ProcessGroupNCCL
|
||||
split_color: int
|
||||
|
||||
def __init__(self, is_high_priority_stream: bool = False): ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
rank: int,
|
||||
size: int,
|
||||
options: Options,
|
||||
) -> None: ...
|
||||
def _group_start(self) -> None: ...
|
||||
def _group_end(self) -> None: ...
|
||||
def _start_time_estimate(self) -> None: ...
|
||||
def _end_time_estimate(self) -> float: ...
|
||||
def _set_default_timeout(self, timeout) -> None: ...
|
||||
def perform_nocolor_split(self, device: torch.device) -> None: ...
|
||||
def register_mem_pool(self, pool: torch.cuda.MemPool) -> None: ...
|
||||
def deregister_mem_pool(self, pool: torch.cuda.MemPool) -> None: ...
|
||||
def comm_split_count(self) -> int: ...
|
||||
def _add_ephemeral_timeout(self, timeout: timedelta) -> None: ...
|
||||
def abort(self) -> None: ...
|
||||
def _is_initialized(self) -> bool: ...
|
||||
@property
|
||||
def uid(self) -> int: ...
|
||||
@property
|
||||
def options(self) -> Options: ... # type: ignore[override]
|
||||
@staticmethod
|
||||
def get_build_nccl_version(self) -> tuple[int, int, int]: ...
|
||||
@staticmethod
|
||||
def get_runtime_nccl_version(self) -> tuple[int, int, int]: ...
|
||||
|
||||
class ProcessGroupUCC(Backend):
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
rank: int,
|
||||
size: int,
|
||||
timeout: timedelta,
|
||||
) -> None: ...
|
||||
|
||||
class ProcessGroupMPI(Backend):
|
||||
def __init__(
|
||||
self,
|
||||
rank: int,
|
||||
size: int,
|
||||
pgComm: int,
|
||||
) -> None: ...
|
||||
@staticmethod
|
||||
def create(ranks: list[int]) -> ProcessGroupMPI: ...
|
||||
|
||||
def _compute_bucket_assignment_by_size(
|
||||
tensors: list[Tensor],
|
||||
bucket_size_limits: list[int],
|
||||
expect_sparse_gradient: list[bool] = ...,
|
||||
tensor_indices: list[int] = ...,
|
||||
) -> tuple[list[list[int]], list[int]]: ...
|
||||
def _broadcast_coalesced(
|
||||
process_group: ProcessGroup,
|
||||
tensors: list[Tensor],
|
||||
buffer_size: int,
|
||||
src: int,
|
||||
): ...
|
||||
def _test_python_store(store: Store): ...
|
||||
def _verify_params_across_processes(
|
||||
process_group: ProcessGroup,
|
||||
params: list[Tensor],
|
||||
logger: Logger | None,
|
||||
): ...
|
||||
def _make_nccl_premul_sum(factor: float | list[Tensor]) -> ReduceOp: ...
|
||||
def _register_process_group(
|
||||
group_name: GroupName,
|
||||
process_group: ProcessGroup,
|
||||
) -> None: ...
|
||||
def _resolve_process_group(group_name: GroupName) -> ProcessGroup: ...
|
||||
def _register_work(tensor: torch.Tensor, work: Work) -> ProcessGroup: ...
|
||||
def _get_work_registry_size() -> int: ...
|
||||
def _set_allow_inflight_collective_as_graph_input(
|
||||
value: bool,
|
||||
) -> None: ...
|
||||
def _allow_inflight_collective_as_graph_input() -> bool: ...
|
||||
def _unregister_all_process_groups() -> None: ...
|
||||
def _unregister_process_group(group_name: GroupName) -> None: ...
|
||||
|
||||
# Initializes the device state in CUmodule so that it's able to perform NVSHMEM
|
||||
# operations. CUmodule is a pointer to a CUDA module, carried by a int64 in
|
||||
# Python. At C++ interface, it is converted to a uintptr_t.
|
||||
def _nvshmemx_cumodule_init(module: int) -> None: ...
|
||||
|
||||
# Check if NVSHMEM is available on current system.
|
||||
def _is_nvshmem_available() -> bool: ...
|
||||
|
||||
class _SymmetricMemory:
|
||||
@staticmethod
|
||||
def set_group_info(
|
||||
group_name: str,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
store: Store,
|
||||
) -> None: ...
|
||||
@staticmethod
|
||||
def empty_strided_p2p(
|
||||
size: torch.types._size,
|
||||
stride: torch.types._size,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
group_name: str | None = None,
|
||||
alloc_id: int | None = None,
|
||||
) -> torch.Tensor: ...
|
||||
@staticmethod
|
||||
def has_multicast_support(
|
||||
device_type: DeviceType,
|
||||
device_idx: int,
|
||||
) -> bool: ...
|
||||
# Set Symmetric Memory allocation backend.
|
||||
@staticmethod
|
||||
def set_backend(name: str) -> None: ...
|
||||
@staticmethod
|
||||
def get_backend(device: torch.device) -> str | None: ...
|
||||
@staticmethod
|
||||
def is_symm_mem_tensor(tensor: torch.Tensor) -> bool: ...
|
||||
@staticmethod
|
||||
def get_mempool_allocator(device: torch.device) -> Any: ...
|
||||
signal_pad_size: int
|
||||
@property
|
||||
def rank(self) -> int: ...
|
||||
@property
|
||||
def world_size(self) -> int: ...
|
||||
@staticmethod
|
||||
def rendezvous(
|
||||
tensor: torch.Tensor, group_name: str | None = None
|
||||
) -> _SymmetricMemory: ...
|
||||
def get_buffer(
|
||||
self,
|
||||
rank: int,
|
||||
sizes: torch.types._size,
|
||||
dtype: torch.dtype,
|
||||
storage_offset: int | None = 0,
|
||||
) -> torch.Tensor: ...
|
||||
def get_signal_pad(
|
||||
self,
|
||||
rank: int,
|
||||
sizes: torch.types._size = [],
|
||||
dtype: torch.dtype | None = None,
|
||||
storage_offset: int | None = 0,
|
||||
) -> torch.Tensor: ...
|
||||
def barrier(self, channel: int = 0, timeout_ms: int = 0) -> None: ...
|
||||
def put_signal(
|
||||
self,
|
||||
dst_rank: int,
|
||||
channel: int = 0,
|
||||
timeout_ms: int = 0,
|
||||
) -> None: ...
|
||||
def wait_signal(
|
||||
self,
|
||||
src_rank: int,
|
||||
channel: int = 0,
|
||||
timeout_ms: int = 0,
|
||||
) -> None: ...
|
||||
def get_remote_tensor(
|
||||
self,
|
||||
peer: int,
|
||||
sizes: torch.types._size,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor: ...
|
||||
@staticmethod
|
||||
def memset32(
|
||||
tensor: torch.Tensor, offset: int, val: int, count: int = 1
|
||||
) -> torch.Tensor: ...
|
||||
@staticmethod
|
||||
def stream_write_value32(
|
||||
tensor: torch.Tensor, offset: int, val: int
|
||||
) -> torch.Tensor: ...
|
||||
@property
|
||||
def buffer_ptrs(self) -> list[int]: ...
|
||||
@property
|
||||
def buffer_ptrs_dev(self) -> int: ...
|
||||
@property
|
||||
def signal_pad_ptrs(self) -> list[int]: ...
|
||||
@property
|
||||
def signal_pad_ptrs_dev(self) -> int: ...
|
||||
@property
|
||||
def multicast_ptr(self) -> int: ...
|
||||
@property
|
||||
def buffer_size(self) -> int: ...
|
||||
@property
|
||||
def device(self) -> torch.device: ...
|
||||
|
||||
class ProcessGroupXCCL(Backend):
|
||||
class Options(Backend.Options):
|
||||
is_high_priority_stream: bool
|
||||
|
||||
def __init__(self, is_high_priority_stream: bool = False): ...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
rank: int,
|
||||
size: int,
|
||||
options: Options,
|
||||
) -> None: ...
|
||||
@property
|
||||
def options(self) -> Options: ... # type: ignore[override]
|
||||
|
||||
def _set_process_group(pg: ProcessGroup) -> None: ...
|
||||
def _current_process_group() -> ProcessGroup: ...
|
||||
|
||||
class _Request:
|
||||
def body(self) -> bytes: ...
|
||||
def get_param(self, str) -> str: ...
|
||||
|
||||
class _Response:
|
||||
def set_content(self, content: str | bytes, content_type: str) -> None: ...
|
||||
def set_status(self, status: int) -> None: ...
|
||||
|
||||
def _register_handler(
|
||||
name: str, handler: Callable[[_Request, _Response], None]
|
||||
) -> None: ...
|
||||
def _set_comm_profiling_name(name: str) -> None: ...
|
||||
def _get_comm_profiling_name() -> str: ...
|
||||
@@ -0,0 +1,188 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# mypy: disable-error-code="type-arg"
|
||||
from datetime import timedelta
|
||||
from typing import Any, Generic, overload, TypeVar
|
||||
|
||||
import torch
|
||||
from torch._C import Future
|
||||
from torch._C._autograd import ProfilerEvent
|
||||
from torch._C._distributed_c10d import Store
|
||||
from torch._C._profiler import ProfilerConfig
|
||||
|
||||
# This module is defined in torch/csrc/distributed/rpc/init.cpp
|
||||
|
||||
_DEFAULT_INIT_METHOD: str
|
||||
_DEFAULT_NUM_WORKER_THREADS: int
|
||||
_UNSET_RPC_TIMEOUT: float
|
||||
_DEFAULT_RPC_TIMEOUT_SEC: float
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
class RpcBackendOptions:
|
||||
rpc_timeout: float
|
||||
init_method: str
|
||||
def __init__(
|
||||
self,
|
||||
rpc_timeout: float = ...,
|
||||
init_method: str = ...,
|
||||
) -> None: ...
|
||||
|
||||
class WorkerInfo:
|
||||
def __init__(self, name: str, worker_id: int) -> None: ...
|
||||
@property
|
||||
def name(self) -> str: ...
|
||||
@property
|
||||
def id(self) -> int: ...
|
||||
def __eq__(self, other: object) -> bool: ...
|
||||
|
||||
class RpcAgent:
|
||||
def join(self, shutdown: bool = False, timeout: float = 0): ...
|
||||
def sync(self): ...
|
||||
def shutdown(self): ...
|
||||
@overload
|
||||
def get_worker_info(self) -> WorkerInfo: ...
|
||||
@overload
|
||||
def get_worker_info(self, workerName: str) -> WorkerInfo: ...
|
||||
def get_worker_infos(self) -> list[WorkerInfo]: ...
|
||||
def _get_device_map(self, dst: WorkerInfo) -> dict[torch.device, torch.device]: ...
|
||||
def get_debug_info(self) -> dict[str, str]: ...
|
||||
def get_metrics(self) -> dict[str, str]: ...
|
||||
|
||||
class PyRRef(Generic[_T]):
|
||||
def __init__(self, value: _T, type_hint: Any = None) -> None: ...
|
||||
def is_owner(self) -> bool: ...
|
||||
def confirmed_by_owner(self) -> bool: ...
|
||||
def owner(self) -> WorkerInfo: ...
|
||||
def owner_name(self) -> str: ...
|
||||
def to_here(self, timeout: float = ...) -> _T: ...
|
||||
def local_value(self) -> Any: ...
|
||||
def rpc_sync(self, timeout: float = ...) -> Any: ...
|
||||
def rpc_async(self, timeout: float = ...) -> Any: ...
|
||||
def remote(self, timeout: float = ...) -> Any: ...
|
||||
def _serialize(self) -> tuple: ...
|
||||
@staticmethod
|
||||
def _deserialize(tp: tuple) -> PyRRef: ...
|
||||
def _get_type(self) -> type[_T]: ...
|
||||
def _get_future(self) -> Future[_T]: ...
|
||||
def _get_profiling_future(self) -> Future[_T]: ...
|
||||
def _set_profiling_future(self, profilingFuture: Future[_T]): ...
|
||||
|
||||
class _TensorPipeRpcBackendOptionsBase(RpcBackendOptions):
|
||||
num_worker_threads: int
|
||||
device_maps: dict[str, dict[torch.device, torch.device]]
|
||||
devices: list[torch.device]
|
||||
def __init__(
|
||||
self,
|
||||
num_worker_threads: int,
|
||||
_transports: list | None,
|
||||
_channels: list | None,
|
||||
rpc_timeout: float = ...,
|
||||
init_method: str = ...,
|
||||
device_maps: dict[str, dict[torch.device, torch.device]] = {}, # noqa: B006
|
||||
devices: list[torch.device] = [], # noqa: B006
|
||||
) -> None: ...
|
||||
def _set_device_map(
|
||||
self,
|
||||
to: str,
|
||||
device_map: dict[torch.device, torch.device],
|
||||
): ...
|
||||
|
||||
class TensorPipeAgent(RpcAgent):
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
name: str,
|
||||
worker_id: int,
|
||||
world_size: int | None,
|
||||
opts: _TensorPipeRpcBackendOptionsBase,
|
||||
reverse_device_maps: dict[str, dict[torch.device, torch.device]],
|
||||
devices: list[torch.device],
|
||||
) -> None: ...
|
||||
def join(self, shutdown: bool = False, timeout: float = 0): ...
|
||||
def shutdown(self): ...
|
||||
@overload
|
||||
def get_worker_info(self) -> WorkerInfo: ...
|
||||
@overload
|
||||
def get_worker_info(self, workerName: str) -> WorkerInfo: ...
|
||||
@overload
|
||||
def get_worker_info(self, id: int) -> WorkerInfo: ...
|
||||
def get_worker_infos(self) -> list[WorkerInfo]: ...
|
||||
def _get_device_map(self, dst: WorkerInfo) -> dict[torch.device, torch.device]: ...
|
||||
def _update_group_membership(
|
||||
self,
|
||||
worker_info: WorkerInfo,
|
||||
my_devices: list[torch.device],
|
||||
reverse_device_map: dict[str, dict[torch.device, torch.device]],
|
||||
is_join: bool,
|
||||
): ...
|
||||
def _get_backend_options(self) -> _TensorPipeRpcBackendOptionsBase: ...
|
||||
@property
|
||||
def is_static_group(self) -> bool: ...
|
||||
@property
|
||||
def store(self) -> Store: ...
|
||||
|
||||
def _is_current_rpc_agent_set() -> bool: ...
|
||||
def _get_current_rpc_agent() -> RpcAgent: ...
|
||||
def _set_and_start_rpc_agent(agent: RpcAgent): ...
|
||||
def _reset_current_rpc_agent(): ...
|
||||
def _delete_all_user_and_unforked_owner_rrefs(timeout: timedelta = ...): ...
|
||||
def _destroy_rref_context(ignoreRRefLeak: bool): ...
|
||||
def _rref_context_get_debug_info() -> dict[str, str]: ...
|
||||
def _cleanup_python_rpc_handler(): ...
|
||||
def _invoke_rpc_builtin(
|
||||
dst: WorkerInfo,
|
||||
opName: str,
|
||||
rpcTimeoutSeconds: float,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
): ...
|
||||
def _invoke_rpc_python_udf(
|
||||
dst: WorkerInfo,
|
||||
pickledPythonUDF: str,
|
||||
tensors: list[torch.Tensor],
|
||||
rpcTimeoutSeconds: float,
|
||||
isAsyncExecution: bool,
|
||||
): ...
|
||||
def _invoke_rpc_torchscript(
|
||||
dstWorkerName: str,
|
||||
qualifiedNameStr: str,
|
||||
argsTuple: tuple,
|
||||
kwargsDict: dict,
|
||||
rpcTimeoutSeconds: float,
|
||||
isAsyncExecution: bool,
|
||||
): ...
|
||||
def _invoke_remote_builtin(
|
||||
dst: WorkerInfo,
|
||||
opName: str,
|
||||
rpcTimeoutSeconds: float,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
): ...
|
||||
def _invoke_remote_python_udf(
|
||||
dst: WorkerInfo,
|
||||
pickledPythonUDF: str,
|
||||
tensors: list[torch.Tensor],
|
||||
rpcTimeoutSeconds: float,
|
||||
isAsyncExecution: bool,
|
||||
): ...
|
||||
def _invoke_remote_torchscript(
|
||||
dstWorkerName: WorkerInfo,
|
||||
qualifiedNameStr: str,
|
||||
rpcTimeoutSeconds: float,
|
||||
isAsyncExecution: bool,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
): ...
|
||||
def get_rpc_timeout() -> float: ...
|
||||
def enable_gil_profiling(flag: bool): ...
|
||||
def _set_rpc_timeout(rpcTimeoutSeconds: float): ...
|
||||
|
||||
class RemoteProfilerManager:
|
||||
@staticmethod
|
||||
def set_current_profiling_key(key: str): ...
|
||||
|
||||
def _enable_server_process_global_profiler(new_config: ProfilerConfig): ...
|
||||
def _disable_server_process_global_profiler() -> list[list[list[ProfilerEvent]]]: ...
|
||||
def _set_profiler_node_id(default_node_id: int): ...
|
||||
def _enable_jit_rref_pickle(): ...
|
||||
def _disable_jit_rref_pickle(): ...
|
||||
@@ -0,0 +1,32 @@
|
||||
import torch
|
||||
from torch._C._distributed_c10d import Store
|
||||
from torch._C._distributed_rpc import _TensorPipeRpcBackendOptionsBase, TensorPipeAgent
|
||||
|
||||
# This module is defined in torch/csrc/distributed/rpc/testing/init.cpp
|
||||
|
||||
class FaultyTensorPipeRpcBackendOptions(_TensorPipeRpcBackendOptionsBase):
|
||||
def __init__(
|
||||
self,
|
||||
num_worker_threads: int,
|
||||
rpc_timeout: float,
|
||||
init_method: str,
|
||||
messages_to_fail: list[str],
|
||||
messages_to_delay: dict[str, float],
|
||||
num_fail_sends: int,
|
||||
) -> None: ...
|
||||
num_send_recv_threads: int
|
||||
messages_to_fail: list[str]
|
||||
messages_to_delay: dict[str, float]
|
||||
num_fail_sends: int
|
||||
|
||||
class FaultyTensorPipeAgent(TensorPipeAgent):
|
||||
def __init__(
|
||||
self,
|
||||
store: Store,
|
||||
name: str,
|
||||
rank: int,
|
||||
world_size: int,
|
||||
options: FaultyTensorPipeRpcBackendOptions,
|
||||
reverse_device_maps: dict[str, dict[torch.device, torch.device]],
|
||||
devices: list[torch.device],
|
||||
) -> None: ...
|
||||
@@ -0,0 +1,69 @@
|
||||
from . import compiled_autograd, eval_frame, guards # noqa: F401
|
||||
|
||||
def strip_function_call(name: str) -> str: ...
|
||||
def is_valid_var_name(name: str) -> bool | int: ...
|
||||
def get_type_slots(obj: type | object) -> tuple[int, int, int, int]: ...
|
||||
def has_slot(slots: int, slot_bit: int) -> bool: ...
|
||||
|
||||
class PySequenceSlots:
|
||||
SQ_LENGTH: int
|
||||
SQ_CONCAT: int
|
||||
SQ_REPEAT: int
|
||||
SQ_ITEM: int
|
||||
SQ_CONTAINS: int
|
||||
SQ_ASS_ITEM: int
|
||||
SQ_INPLACE_CONCAT: int
|
||||
SQ_INPLACE_REPEAT: int
|
||||
|
||||
class PyMappingSlots:
|
||||
MP_LENGTH: int
|
||||
MP_SUBSCRIPT: int
|
||||
MP_ASS_SUBSCRIPT: int
|
||||
|
||||
class PyNumberSlots:
|
||||
NB_ADD: int
|
||||
NB_SUBTRACT: int
|
||||
NB_MULTIPLY: int
|
||||
NB_REMAINDER: int
|
||||
NB_POWER: int
|
||||
NB_NEGATIVE: int
|
||||
NB_POSITIVE: int
|
||||
NB_ABSOLUTE: int
|
||||
NB_BOOL: int
|
||||
NB_INVERT: int
|
||||
NB_LSHIFT: int
|
||||
NB_RSHIFT: int
|
||||
NB_AND: int
|
||||
NB_XOR: int
|
||||
NB_OR: int
|
||||
NB_INT: int
|
||||
NB_FLOAT: int
|
||||
NB_INPLACE_ADD: int
|
||||
NB_INPLACE_SUBTRACT: int
|
||||
NB_INPLACE_MULTIPLY: int
|
||||
NB_INPLACE_REMAINDER: int
|
||||
NB_INPLACE_POWER: int
|
||||
NB_INPLACE_LSHIFT: int
|
||||
NB_INPLACE_RSHIFT: int
|
||||
NB_INPLACE_AND: int
|
||||
NB_INPLACE_XOR: int
|
||||
NB_INPLACE_OR: int
|
||||
NB_FLOOR_DIVIDE: int
|
||||
NB_TRUE_DIVIDE: int
|
||||
NB_INPLACE_FLOOR_DIVIDE: int
|
||||
NB_INPLACE_TRUE_DIVIDE: int
|
||||
NB_INDEX: int
|
||||
NB_MATRIX_MULTIPLY: int
|
||||
NB_INPLACE_MATRIX_MULTIPLY: int
|
||||
|
||||
class PyTypeSlots:
|
||||
TP_HASH: int
|
||||
TP_ITER: int
|
||||
TP_ITERNEXT: int
|
||||
TP_CALL: int
|
||||
TP_REPR: int
|
||||
TP_RICHCOMPARE: int
|
||||
TP_GETATTRO: int
|
||||
TP_SETATTRO: int
|
||||
TP_DESCR_GET: int
|
||||
TP_DESCR_SET: int
|
||||
@@ -0,0 +1,13 @@
|
||||
from collections.abc import Callable
|
||||
|
||||
from torch import Tensor
|
||||
from torch._dynamo.compiled_autograd import AutogradCompilerInstance
|
||||
|
||||
def set_autograd_compiler(
|
||||
autograd_compiler: Callable[[], AutogradCompilerInstance] | None,
|
||||
dynamic: bool,
|
||||
) -> tuple[Callable[[], AutogradCompilerInstance] | None, bool]: ...
|
||||
def clear_cache() -> None: ...
|
||||
def is_cache_empty() -> bool: ...
|
||||
def set_verbose_logger(fn: Callable[[str], None] | None) -> bool: ...
|
||||
def call_cpp_tensor_pre_hooks(idx: int, grad: Tensor) -> Tensor: ...
|
||||
@@ -0,0 +1,99 @@
|
||||
import enum
|
||||
import types
|
||||
from collections.abc import Callable
|
||||
from typing import overload
|
||||
|
||||
from torch._dynamo.guards import GuardManagerWrapper
|
||||
from torch._dynamo.types import DynamoCallback, DynamoGuardCompleteHook, DynamoGuardHook
|
||||
from torch._guards import CompileId
|
||||
|
||||
def set_eval_frame(callback: DynamoCallback) -> DynamoCallback: ...
|
||||
def set_skip_guard_eval_unsafe(value: bool) -> bool: ...
|
||||
def get_eval_frame_callback() -> DynamoCallback: ...
|
||||
def reset_code(code: types.CodeType) -> None: ...
|
||||
def unsupported(obj1: object, obj2: object) -> object: ...
|
||||
def set_code_exec_strategy(
|
||||
code: types.CodeType, strategy: _FrameExecStrategy
|
||||
) -> None: ...
|
||||
def set_guard_error_hook(hook: DynamoGuardHook) -> None: ...
|
||||
def set_guard_complete_hook(
|
||||
hook: DynamoGuardCompleteHook | None,
|
||||
) -> DynamoGuardCompleteHook | None: ...
|
||||
def raise_sigtrap() -> None: ...
|
||||
def set_c_recursion_limit(limit: int) -> None: ...
|
||||
def get_c_recursion_limit() -> int: ...
|
||||
|
||||
class _CacheEntry:
|
||||
def check_fn(self, *args: object, **kwargs: object) -> bool: ...
|
||||
def update_diff_guard_root_manager(self) -> None: ...
|
||||
code: types.CodeType
|
||||
compile_id: CompileId
|
||||
# If we run into circular issues, just use object
|
||||
guard_manager: GuardManagerWrapper
|
||||
backend: Callable
|
||||
next: _CacheEntry | None
|
||||
|
||||
class _PrecompileEntry:
|
||||
guard_manager: GuardManagerWrapper
|
||||
|
||||
class _ExtraState:
|
||||
def invalidate(
|
||||
self, cache_entry: _CacheEntry, guard_manager: GuardManagerWrapper
|
||||
) -> None: ...
|
||||
|
||||
class _FrameAction(enum.IntEnum):
|
||||
DEFAULT = 0
|
||||
SKIP = 1
|
||||
RUN_ONLY = 2
|
||||
|
||||
class _FrameExecStrategy:
|
||||
cur_action: _FrameAction
|
||||
recursive_action: _FrameAction
|
||||
|
||||
@overload
|
||||
def __init__(self) -> None: ...
|
||||
@overload
|
||||
def __init__(
|
||||
self, cur_action: _FrameAction, recursive_action: _FrameAction
|
||||
) -> None: ...
|
||||
|
||||
# This is an object that encapsulates the Python FrameType, and exposes
|
||||
# properties Dynamo cares about for a frame.
|
||||
class _PyInterpreterFrame:
|
||||
f_code: types.CodeType
|
||||
f_locals: dict[str, object]
|
||||
f_globals: dict[str, object]
|
||||
f_builtins: dict[str, object]
|
||||
f_lasti: int
|
||||
f_lineno: int
|
||||
f_back: types.FrameType
|
||||
# A tuple containing cell objects captured by this frame.
|
||||
closure: tuple[types.CellType]
|
||||
|
||||
def _debug_get_cache_entry_list(code: types.CodeType) -> list[_CacheEntry]: ...
|
||||
def _get_frame_value_stack_with_depth(
|
||||
frame: types.FrameType, depth: int
|
||||
) -> list[object]: ...
|
||||
def set_bytecode_debugger_callback(
|
||||
callback: Callable[[types.CodeType], None] | None,
|
||||
) -> None: ...
|
||||
def get_bytecode_debugger_callback() -> Callable[[types.CodeType], None] | None: ...
|
||||
def register_breakpoint_code(code: types.CodeType) -> None: ...
|
||||
|
||||
# Sentinel for NULL stack values returned by _get_frame_value_stack_at_depth
|
||||
class NullStackValue: ...
|
||||
|
||||
NULL_STACK_VALUE: NullStackValue
|
||||
|
||||
py_opcode_caches: list[int]
|
||||
|
||||
def code_framelocals_names(code: types.CodeType) -> tuple[str, ...]: ...
|
||||
def _load_precompile_entry(
|
||||
code: types.CodeType,
|
||||
guard_manager: GuardManagerWrapper,
|
||||
dynamo_code: types.CodeType,
|
||||
) -> None: ...
|
||||
def _reset_precompile_entries(code: types.CodeType) -> None: ...
|
||||
def _debug_get_precompile_entries(code: types.CodeType) -> list[_PrecompileEntry]: ...
|
||||
def set_fullgraph_compiled_frame_count(value: int) -> int: ...
|
||||
def set_fullgraph_error_on_nested_compile(value: bool) -> bool: ...
|
||||
@@ -0,0 +1,500 @@
|
||||
import enum
|
||||
import traceback
|
||||
from collections.abc import Callable
|
||||
from typing import Any, TypeAlias
|
||||
|
||||
import torch
|
||||
|
||||
# TODO: We should move the `GuardManagerType`
|
||||
# defined in `guards.py` here and update other
|
||||
# imports
|
||||
GuardManagerType: TypeAlias = enum.Enum
|
||||
|
||||
class GlobalStateGuard:
|
||||
def check(self) -> bool: ...
|
||||
def reason(self) -> str: ...
|
||||
|
||||
class LeafGuard:
|
||||
def verbose_code_parts(self) -> list[str]: ...
|
||||
|
||||
class RelationalGuard: ...
|
||||
|
||||
class GuardDebugInfo:
|
||||
verbose_code_parts: list[str]
|
||||
result: bool
|
||||
num_guards_executed: int
|
||||
user_stack: traceback.StackSummary | None
|
||||
|
||||
class GuardManager:
|
||||
def check(self, value: Any) -> bool: ...
|
||||
def check_verbose(self, value: Any) -> GuardDebugInfo: ...
|
||||
|
||||
# Accessors
|
||||
def globals_dict_manager(
|
||||
self,
|
||||
f_globals: dict[str, Any],
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def framelocals_manager(
|
||||
self,
|
||||
key: tuple[str, int],
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def dict_getitem_manager(
|
||||
self,
|
||||
key: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def grad_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def generic_getattr_manager(
|
||||
self,
|
||||
attr: str,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def getitem_manager(
|
||||
self,
|
||||
key: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def get_generic_dict_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def list_getitem_manager(
|
||||
self,
|
||||
key: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def tuple_getitem_manager(
|
||||
self,
|
||||
key: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def set_getitem_manager(
|
||||
self,
|
||||
index: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def func_defaults_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def func_kwdefaults_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def tuple_iterator_getitem_manager(
|
||||
self,
|
||||
index: Any,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def weakref_call_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def call_function_no_args_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def global_weakref_manager(
|
||||
self,
|
||||
global_name: str,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def type_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def getattr_manager(
|
||||
self,
|
||||
attr: str,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def tensor_property_size_manager(
|
||||
self,
|
||||
idx: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def tensor_property_shape_manager(
|
||||
self,
|
||||
idx: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def tensor_property_storage_offset_manager(
|
||||
self,
|
||||
idx: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def indexed_manager(
|
||||
self,
|
||||
idx: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def lambda_manager(
|
||||
self,
|
||||
python_lambda: Callable[..., Any],
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def get_root(self) -> RootGuardManager: ...
|
||||
def get_source(self) -> str: ...
|
||||
def fail_count(self) -> int: ...
|
||||
def get_child_managers(self) -> list[GuardManager]: ...
|
||||
def repr(self) -> str: ...
|
||||
def type_of_guarded_value(self) -> str: ...
|
||||
def get_leaf_guards(self) -> list[LeafGuard]: ...
|
||||
def get_accessors(self) -> list[GuardManager]: ...
|
||||
def is_guarded_value_immutable(self) -> bool: ...
|
||||
def is_tag_safe(self) -> bool: ...
|
||||
def is_tag_safe_root(self) -> bool: ...
|
||||
def has_no_accessors(self) -> bool: ...
|
||||
def has_object_aliasing_guard(self) -> bool: ...
|
||||
def get_type_of_guarded_value(self) -> type: ...
|
||||
def type_dict_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def type_mro_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def code_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def closure_manager(
|
||||
self,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
# Leaf guards
|
||||
def add_lambda_guard(
|
||||
self,
|
||||
user_lambda: Callable[..., Any],
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_id_match_guard(
|
||||
self,
|
||||
id_val: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_equals_match_guard(
|
||||
self,
|
||||
equals_val: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_global_state_guard(
|
||||
self,
|
||||
initial_state: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_torch_function_mode_stack_guard(
|
||||
self,
|
||||
initial_stack: list[Any],
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_mapping_keys_guard(
|
||||
self,
|
||||
value: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_dict_length_check_guard(
|
||||
self,
|
||||
value: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_length_check_guard(
|
||||
self,
|
||||
value: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_true_match_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_false_match_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_none_match_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_not_none_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_dispatch_key_set_guard(
|
||||
self,
|
||||
dispatch_key: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_tensor_match_guard(
|
||||
self,
|
||||
value: Any,
|
||||
sizes: list[int],
|
||||
strides: list[int],
|
||||
tensor_name: str,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
ptype: Any,
|
||||
dispatch_keys: Any,
|
||||
) -> None: ...
|
||||
def add_dynamic_indices_guard(
|
||||
self,
|
||||
value: set[Any],
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_no_hasattr_guard(
|
||||
self,
|
||||
attr_name: str,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_dict_contains_guard(
|
||||
self,
|
||||
contains: bool,
|
||||
key: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_type_match_guard(
|
||||
self,
|
||||
value: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_dict_version_guard(
|
||||
self,
|
||||
value: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_set_contains_guard(
|
||||
self,
|
||||
contains: bool,
|
||||
item: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_dual_level_match_guard(
|
||||
self,
|
||||
level: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_float_is_nan_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_complex_is_nan_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_tuple_iterator_length_guard(
|
||||
self,
|
||||
length: int,
|
||||
type_id: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_range_iterator_match_guard(
|
||||
self,
|
||||
start: int,
|
||||
stop: int,
|
||||
step: int,
|
||||
type_id: int,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def add_default_device_guard(
|
||||
self,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def mark_tag_safe(self) -> None: ...
|
||||
def mark_tag_safe_root(self) -> None: ...
|
||||
|
||||
class RootGuardManager(GuardManager):
|
||||
def get_epilogue_lambda_guards(self) -> list[LeafGuard]: ...
|
||||
def add_epilogue_lambda_guard(
|
||||
self,
|
||||
guard: LeafGuard,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def clone_manager(
|
||||
self, clone_filter_fn: Callable[[GuardManager], bool]
|
||||
) -> RootGuardManager: ...
|
||||
def attach_compile_id(self, compile_id: str) -> None: ...
|
||||
|
||||
class DictGuardManager(GuardManager):
|
||||
def get_key_manager(
|
||||
self,
|
||||
index: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def get_value_manager(
|
||||
self,
|
||||
index: int,
|
||||
source: str,
|
||||
example_value: Any,
|
||||
guard_manager_enum: GuardManagerType,
|
||||
) -> GuardManager: ...
|
||||
def get_key_value_managers(
|
||||
self,
|
||||
) -> dict[int, tuple[GuardManager, GuardManager]]: ...
|
||||
|
||||
# Guard accessor stubs
|
||||
class GuardAccessor: ...
|
||||
class DictGetItemGuardAccessor(GuardAccessor): ...
|
||||
class GetGenericDictGuardAccessor(GuardAccessor): ...
|
||||
class TypeDictGuardAccessor(GuardAccessor): ...
|
||||
class TypeMROGuardAccessor(GuardAccessor): ...
|
||||
class ClosureGuardAccessor(GuardAccessor): ...
|
||||
class TupleGetItemGuardAccessor(GuardAccessor): ...
|
||||
class TypeGuardAccessor(GuardAccessor): ...
|
||||
class CodeGuardAccessor(GuardAccessor): ...
|
||||
class FuncDefaultsGuardAccessor(GuardAccessor): ...
|
||||
class FuncKwDefaultsGuardAccessor(GuardAccessor): ...
|
||||
|
||||
class GetAttrGuardAccessor(GuardAccessor):
|
||||
def get_attr_name(self) -> str: ...
|
||||
|
||||
def install_object_aliasing_guard(
|
||||
x: GuardManager,
|
||||
y: GuardManager,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def install_no_tensor_aliasing_guard(
|
||||
guard_managers: list[GuardManager],
|
||||
tensor_names: list[str],
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def install_storage_overlapping_guard(
|
||||
overlapping_guard_managers: list[GuardManager],
|
||||
non_overlapping_guard_managers: list[GuardManager],
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def install_symbolic_shape_guard(
|
||||
guard_managers: list[GuardManager],
|
||||
nargs_int: int,
|
||||
nargs_float: int,
|
||||
py_addr: int,
|
||||
py_addr_keep_alive: Any,
|
||||
verbose_code_parts: list[str],
|
||||
user_stack: traceback.StackSummary | None,
|
||||
) -> None: ...
|
||||
def profile_guard_manager(
|
||||
guard_manager: GuardManager,
|
||||
f_locals: dict[str, Any],
|
||||
n_iters: int,
|
||||
) -> float: ...
|
||||
|
||||
class TensorGuards:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dynamic_dims_sizes: list[torch.SymInt | None] | None = None,
|
||||
dynamic_dims_strides: list[torch.SymInt | None] | None = None,
|
||||
) -> None: ...
|
||||
def check(self, *args: Any) -> bool: ...
|
||||
def check_verbose(
|
||||
self, *args: Any, tensor_check_names: list[str] | None = None
|
||||
) -> bool | str: ...
|
||||
|
||||
def assert_size_stride(
|
||||
item: torch.Tensor,
|
||||
size: torch.types._size,
|
||||
stride: torch.types._size,
|
||||
op_name: str | None = None,
|
||||
) -> None: ...
|
||||
def assert_alignment(
|
||||
item: torch.Tensor,
|
||||
alignment: int,
|
||||
op_name: str | None = None,
|
||||
) -> None: ...
|
||||
def copy_misaligned(item: torch.Tensor) -> torch.Tensor: ...
|
||||
def check_obj_id(obj: object, expected: int) -> bool: ...
|
||||
def check_type_id(obj: object, expected: int) -> bool: ...
|
||||
def dict_version(d: dict[Any, Any]) -> int: ...
|
||||
def compute_overlapping_tensors(
|
||||
tensors: list[torch.Tensor], symbolic: bool = True
|
||||
) -> set[int]: ...
|
||||
def set_is_in_mode_without_ignore_compile_internals(value: bool) -> None: ...
|
||||
@@ -0,0 +1,9 @@
|
||||
# Defined in torch/csrc/export/pybind.cpp
|
||||
class CppExportedProgram: ...
|
||||
|
||||
def deserialize_exported_program(
|
||||
serialized_program: str,
|
||||
) -> CppExportedProgram: ...
|
||||
def serialize_exported_program(
|
||||
cpp_exported_program: CppExportedProgram,
|
||||
) -> str: ...
|
||||
@@ -0,0 +1,25 @@
|
||||
# Defined in torch/csrc/export/pt2_archive_constants.h
|
||||
|
||||
ARCHIVE_ROOT_NAME: str = ...
|
||||
ARCHIVE_FORMAT_PATH: str = ...
|
||||
ARCHIVE_FORMAT_VALUE: str = ...
|
||||
ARCHIVE_VERSION_PATH: str = ...
|
||||
ARCHIVE_VERSION_VALUE: str = ...
|
||||
MODELS_DIR: str = ...
|
||||
MODELS_FILENAME_FORMAT: str = ...
|
||||
AOTINDUCTOR_DIR: str = ...
|
||||
MTIA_DIR: str = ...
|
||||
WEIGHTS_DIR: str = ...
|
||||
WEIGHTS_CONFIG_FILENAME_FORMAT: str = ...
|
||||
WEIGHT_FILENAME_PREFIX: str = ...
|
||||
CONSTANTS_DIR: str = ...
|
||||
CONSTANTS_CONFIG_FILENAME_FORMAT: str = ...
|
||||
TENSOR_CONSTANT_FILENAME_PREFIX: str = ...
|
||||
CUSTOM_OBJ_FILENAME_PREFIX: str = ...
|
||||
SAMPLE_INPUTS_DIR: str = ...
|
||||
SAMPLE_INPUTS_FILENAME_FORMAT: str = ...
|
||||
EXECUTORCH_DIR: str = ...
|
||||
EXTRA_DIR: str = ...
|
||||
MODULE_INFO_PATH: str = ...
|
||||
XL_MODEL_WEIGHTS_DIR: str = ...
|
||||
XL_MODEL_WEIGHTS_PARAM_CONFIG_PATH: str = ...
|
||||
@@ -0,0 +1,16 @@
|
||||
from torch import Tensor
|
||||
from torch.types import _bool
|
||||
|
||||
# Defined in torch/csrc/functionalization/Module.cpp
|
||||
|
||||
class ViewMeta:
|
||||
has_symbolic_inputs: _bool
|
||||
|
||||
# Returns the list of ViewMeta instances of the given functional tensor.
|
||||
#
|
||||
# Although we do have python bindings for their types, we won't
|
||||
# expose them here, since they should not be used by users.
|
||||
def get_view_meta_sequence(tensor: Tensor) -> list[ViewMeta]: ...
|
||||
|
||||
# Applies the ViewMeta sequence on top of the given base.
|
||||
def apply_view_meta_sequence(base: Tensor, sequence: list[ViewMeta]) -> Tensor: ...
|
||||
@@ -0,0 +1,19 @@
|
||||
from typing import AnyStr, overload
|
||||
|
||||
from torch import Tensor
|
||||
|
||||
class UndefinedGrad:
|
||||
def __init__(self) -> None: ...
|
||||
def __call__(self, *inputs: Tensor) -> list[Tensor]: ...
|
||||
|
||||
class DelayedError:
|
||||
def __init__(self, msg: AnyStr, num_inputs: int) -> None: ...
|
||||
|
||||
# __call__ should really be a higher-kinded type:
|
||||
# def __call__(self, arg: Tensor) -> Tensor: ...
|
||||
# def __call__(self, *args: Tensor * num_inputs) -> Tuple[Tensor * num_inputs]: ...
|
||||
|
||||
@overload
|
||||
def __call__(self, i0: Tensor) -> Tensor: ...
|
||||
@overload
|
||||
def __call__(self, *args: Tensor) -> tuple[Tensor, ...]: ...
|
||||
@@ -0,0 +1,100 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from enum import Enum
|
||||
|
||||
from torch import Tensor
|
||||
|
||||
# Defined in torch/csrc/functorch/init.cpp
|
||||
|
||||
def _assert_wrapped_functional(
|
||||
input_tensor: Tensor, wrapped_tensor: Tensor
|
||||
) -> None: ...
|
||||
def _func_decrement_nesting() -> int: ...
|
||||
def _func_increment_nesting(reapply_views: bool) -> int: ...
|
||||
def _propagate_functional_input_mutation(
|
||||
input_tensor: Tensor, wrapped_tensor: Tensor
|
||||
) -> None: ...
|
||||
def set_inplace_requires_grad_allowed(allowed: bool) -> None: ...
|
||||
def get_inplace_requires_grad_allowed() -> bool: ...
|
||||
def _set_dynamic_layer_keys_included(included: bool) -> None: ...
|
||||
def get_unwrapped(tensor: Tensor) -> Tensor: ...
|
||||
def is_batchedtensor(tensor: Tensor) -> bool: ...
|
||||
def is_functionaltensor(tensor: Tensor) -> bool: ...
|
||||
def is_functorch_wrapped_tensor(tensor: Tensor) -> bool: ...
|
||||
def is_gradtrackingtensor(tensor: Tensor) -> bool: ...
|
||||
def is_legacy_batchedtensor(tensor: Tensor) -> bool: ...
|
||||
def maybe_get_bdim(tensor: Tensor) -> int: ...
|
||||
def maybe_get_level(tensor: Tensor) -> int: ...
|
||||
def maybe_current_level() -> int | None: ...
|
||||
def unwrap_if_dead(tensor: Tensor) -> Tensor: ...
|
||||
def _unwrap_for_grad(tensor: Tensor, level: int) -> Tensor: ...
|
||||
def _wrap_for_grad(tensor: Tensor, level: int) -> Tensor: ...
|
||||
def _unwrap_batched(tensor: Tensor, level: int) -> tuple[Tensor, int | None]: ...
|
||||
def current_level() -> int: ...
|
||||
def count_jvp_interpreters() -> int: ...
|
||||
def _add_batch_dim(tensor: Tensor, bdim: int, level: int) -> Tensor: ...
|
||||
def _remove_batch_dim(
|
||||
tensor: Tensor, level: int, batch_size: int, out_dim: int
|
||||
) -> Tensor: ...
|
||||
def _maybe_unsafe_set_level(tensor: Tensor, level: int) -> None: ...
|
||||
def set_single_level_autograd_function_allowed(allowed: bool) -> None: ...
|
||||
def get_single_level_autograd_function_allowed() -> bool: ...
|
||||
def _unwrap_functional_tensor(tensor: Tensor, reapply_views: bool) -> Tensor: ...
|
||||
def _wrap_functional_tensor(tensor: Tensor, level: int) -> Tensor: ...
|
||||
def _vmap_increment_nesting(batch_size: int, randomness: str) -> int: ...
|
||||
def _vmap_decrement_nesting() -> int: ...
|
||||
def _grad_increment_nesting() -> int: ...
|
||||
def _grad_decrement_nesting() -> int: ...
|
||||
def _jvp_increment_nesting() -> int: ...
|
||||
def _jvp_decrement_nesting() -> int: ...
|
||||
|
||||
# Defined in aten/src/ATen/functorch/Interpreter.h
|
||||
class TransformType(Enum):
|
||||
Torch = ...
|
||||
Vmap = ...
|
||||
Grad = ...
|
||||
Jvp = ...
|
||||
Functionalize = ...
|
||||
|
||||
class RandomnessType(Enum):
|
||||
Error = ...
|
||||
Same = ...
|
||||
Different = ...
|
||||
|
||||
class CInterpreter:
|
||||
def key(self) -> TransformType: ...
|
||||
def level(self) -> int: ...
|
||||
def serialize(self) -> bytes: ...
|
||||
@staticmethod
|
||||
def deserialize(bytes) -> CInterpreter: ...
|
||||
|
||||
class CGradInterpreterPtr:
|
||||
def __init__(self, interpreter: CInterpreter) -> None: ...
|
||||
def lift(self, Tensor) -> Tensor: ...
|
||||
def prevGradMode(self) -> bool: ...
|
||||
|
||||
class CJvpInterpreterPtr:
|
||||
def __init__(self, interpreter: CInterpreter) -> None: ...
|
||||
def lift(self, Tensor) -> Tensor: ...
|
||||
def prevFwdGradMode(self) -> bool: ...
|
||||
|
||||
class CFunctionalizeInterpreterPtr:
|
||||
def __init__(self, interpreter: CInterpreter) -> None: ...
|
||||
def key(self) -> TransformType: ...
|
||||
def level(self) -> int: ...
|
||||
def functionalizeAddBackViews(self) -> bool: ...
|
||||
|
||||
class CVmapInterpreterPtr:
|
||||
def __init__(self, interpreter: CInterpreter) -> None: ...
|
||||
def key(self) -> TransformType: ...
|
||||
def level(self) -> int: ...
|
||||
def batchSize(self) -> int: ...
|
||||
def randomness(self) -> RandomnessType: ...
|
||||
|
||||
class DynamicLayer: ...
|
||||
|
||||
def get_dynamic_layer_stack_depth() -> int: ...
|
||||
def get_interpreter_stack() -> list[CInterpreter]: ...
|
||||
def peek_interpreter_stack() -> CInterpreter: ...
|
||||
def pop_dynamic_layer_stack() -> DynamicLayer: ...
|
||||
def pop_dynamic_layer_stack_and_undo_to_depth(int) -> None: ...
|
||||
def push_dynamic_layer_stack(dl: DynamicLayer) -> int: ...
|
||||
@@ -0,0 +1,4 @@
|
||||
# Defined in torch/csrc/instruction_counter/Module.cpp
|
||||
|
||||
def start() -> int: ...
|
||||
def end(id: int) -> int: ...
|
||||
@@ -0,0 +1,5 @@
|
||||
# Defined in torch/csrc/itt.cpp
|
||||
def is_available() -> None: ...
|
||||
def rangePush(message: str) -> None: ...
|
||||
def rangePop() -> None: ...
|
||||
def mark(message: str) -> None: ...
|
||||
@@ -0,0 +1,200 @@
|
||||
from typing import Any
|
||||
|
||||
# Defined in torch/csrc/jit/python/python_tree_views.cpp
|
||||
|
||||
class SourceRange:
|
||||
def highlight(self) -> str: ...
|
||||
@property
|
||||
def start(self) -> int: ...
|
||||
@property
|
||||
def end(self) -> int: ...
|
||||
|
||||
class SourceRangeFactory:
|
||||
def __init__(
|
||||
self,
|
||||
text: str,
|
||||
filename: Any,
|
||||
file_lineno: int,
|
||||
leading_whitespace_chars: int,
|
||||
) -> None: ...
|
||||
def make_range(self, line: int, start_col: int, end_col: int) -> SourceRange: ...
|
||||
def make_raw_range(self, start: int, end: int) -> SourceRange: ...
|
||||
@property
|
||||
def source(self) -> str: ...
|
||||
|
||||
class TreeView:
|
||||
def range(self) -> SourceRange: ...
|
||||
def dump(self) -> None: ...
|
||||
|
||||
class Ident(TreeView):
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None: ...
|
||||
@property
|
||||
def name(self) -> str: ...
|
||||
|
||||
class Param(TreeView):
|
||||
def __init__(self, type: Any | None, name: Ident, kwarg_only: bool) -> None: ...
|
||||
|
||||
class Attribute(TreeView):
|
||||
def __init__(self, name: Ident, value: Any) -> None: ...
|
||||
|
||||
# Literals
|
||||
def TrueLiteral(range: SourceRange) -> Any: ...
|
||||
def FalseLiteral(range: SourceRange) -> Any: ...
|
||||
def NoneLiteral(range: SourceRange) -> Any: ...
|
||||
|
||||
# Tree nodes
|
||||
class Stmt(TreeView):
|
||||
def __init__(self, thing: TreeView) -> None: ...
|
||||
|
||||
class Expr(TreeView): ...
|
||||
|
||||
class Def(TreeView):
|
||||
def __init__(self, name: Ident, decl: Any, body: list[Stmt]) -> None: ...
|
||||
def decl(self) -> Any: ...
|
||||
def name(self) -> Ident: ...
|
||||
|
||||
class Property(TreeView):
|
||||
def __init__(
|
||||
self, r: SourceRange, name: Ident, getter: Def, setter: Def | None
|
||||
) -> None: ...
|
||||
def name(self) -> Ident: ...
|
||||
def getter_name(self) -> str: ...
|
||||
def setter_name(self) -> Ident | None: ...
|
||||
|
||||
class ClassDef(TreeView):
|
||||
def __init__(
|
||||
self, name: Ident, body: list[Stmt], props: list[Property], assigns: list[Any]
|
||||
) -> None: ...
|
||||
|
||||
class Decl(TreeView):
|
||||
def __init__(
|
||||
self, r: SourceRange, params: list[Param], return_type: Expr | None
|
||||
) -> None: ...
|
||||
|
||||
class Delete(Stmt):
|
||||
def __init__(self, range: SourceRange, targets: list[Expr]) -> None: ...
|
||||
|
||||
class WithItem(Expr):
|
||||
def __init__(self, range: SourceRange, target: Expr, var: Any | None) -> None: ...
|
||||
|
||||
class Assign(Stmt):
|
||||
def __init__(
|
||||
self, lhs: list[Expr], rhs: Expr, type: Expr | None = None
|
||||
) -> None: ...
|
||||
|
||||
class AugAssign(Stmt):
|
||||
def __init__(self, lhs: Expr, kind_str: str, rhs: Expr) -> None: ...
|
||||
|
||||
class Return(Stmt):
|
||||
def __init__(self, range: SourceRange, value: Expr | None) -> None: ...
|
||||
|
||||
class Raise(Stmt):
|
||||
def __init__(self, range: SourceRange, expr: Expr) -> None: ...
|
||||
|
||||
class Assert(Stmt):
|
||||
def __init__(self, range: SourceRange, test: Expr, msg: Expr | None) -> None: ...
|
||||
|
||||
class Pass(Stmt):
|
||||
def __init__(self, range: SourceRange) -> None: ...
|
||||
|
||||
class Break(Stmt): ...
|
||||
class Continue(Stmt): ...
|
||||
|
||||
class Dots(Expr, TreeView):
|
||||
def __init__(self, range: SourceRange) -> None: ...
|
||||
|
||||
class If(Stmt):
|
||||
def __init__(
|
||||
self,
|
||||
range: SourceRange,
|
||||
cond: Expr,
|
||||
true_branch: list[Stmt],
|
||||
false_branch: list[Stmt],
|
||||
) -> None: ...
|
||||
|
||||
class While(Stmt):
|
||||
def __init__(self, range: SourceRange, cond: Expr, body: list[Stmt]) -> None: ...
|
||||
|
||||
class With(Stmt):
|
||||
def __init__(
|
||||
self, range: SourceRange, targets: list[WithItem], body: list[Stmt]
|
||||
) -> None: ...
|
||||
|
||||
class For(Stmt):
|
||||
def __init__(
|
||||
self,
|
||||
range: SourceRange,
|
||||
targets: list[Expr],
|
||||
itrs: list[Expr],
|
||||
body: list[Stmt],
|
||||
) -> None: ...
|
||||
|
||||
class ExprStmt(Stmt):
|
||||
def __init__(self, expr: Expr) -> None: ...
|
||||
|
||||
class Var(Expr):
|
||||
def __init__(self, name: Ident) -> None: ...
|
||||
@property
|
||||
def name(self) -> str: ...
|
||||
|
||||
class BinOp(Expr):
|
||||
def __init__(self, kind: str, lhs: Expr, rhs: Expr) -> None: ...
|
||||
|
||||
class UnaryOp(Expr):
|
||||
def __init__(self, range: SourceRange, kind: str, expr: Expr) -> None: ...
|
||||
|
||||
class Const(Expr):
|
||||
def __init__(self, range: SourceRange, value: str) -> None: ...
|
||||
|
||||
class StringLiteral(Expr):
|
||||
def __init__(self, range: SourceRange, value: str) -> None: ...
|
||||
|
||||
class Apply(Expr):
|
||||
def __init__(
|
||||
self, expr: Expr, args: list[Expr], kwargs: list[Attribute]
|
||||
) -> None: ...
|
||||
|
||||
class Select(Expr):
|
||||
def __init__(self, expr: Expr, field: Ident) -> None: ...
|
||||
|
||||
class TernaryIf(Expr):
|
||||
def __init__(self, cond: Expr, true_expr: Expr, false_expr: Expr) -> None: ...
|
||||
|
||||
class ListComp(Expr):
|
||||
def __init__(
|
||||
self, range: SourceRange, elt: Expr, target: Expr, iter: Expr
|
||||
) -> None: ...
|
||||
|
||||
class DictComp(Expr):
|
||||
def __init__(
|
||||
self, range: SourceRange, key: Expr, value: Expr, target: Expr, iter: Expr
|
||||
) -> None: ...
|
||||
|
||||
class ListLiteral(Expr):
|
||||
def __init__(self, range: SourceRange, args: list[Expr]) -> None: ...
|
||||
|
||||
class TupleLiteral(Expr):
|
||||
def __init__(self, range: SourceRange, args: list[Expr]) -> None: ...
|
||||
|
||||
class DictLiteral(Expr):
|
||||
def __init__(
|
||||
self, range: SourceRange, keys: list[Expr], values: list[Expr]
|
||||
) -> None: ...
|
||||
|
||||
class Subscript(Expr):
|
||||
def __init__(self, base: Expr, subscript_exprs: list[Expr]) -> None: ...
|
||||
|
||||
class SliceExpr(Expr):
|
||||
def __init__(
|
||||
self,
|
||||
range: SourceRange,
|
||||
lower: Expr | None,
|
||||
upper: Expr | None,
|
||||
step: Expr | None,
|
||||
) -> None: ...
|
||||
|
||||
class Starred(Expr):
|
||||
def __init__(self, range: SourceRange, expr: Expr) -> None: ...
|
||||
|
||||
class EmptyTypeAnnotation(TreeView):
|
||||
def __init__(self, range: SourceRange) -> None: ...
|
||||
@@ -0,0 +1,26 @@
|
||||
from torch import Tensor
|
||||
|
||||
# defined in torch/csrc/lazy/python/init.cpp
|
||||
def _mark_step(device: str, devices: list[str], wait: bool) -> None: ...
|
||||
def _wait_device_ops(devices: list[str]) -> None: ...
|
||||
def _reset_metrics() -> None: ...
|
||||
def _counter_names() -> list[str]: ...
|
||||
def _counter_value(name: str) -> int: ...
|
||||
def _metrics_report() -> str: ...
|
||||
def _get_graph_hash(tensors: list[Tensor]) -> str: ...
|
||||
def _sync_multi(
|
||||
tensors: list[Tensor],
|
||||
devices: list[str],
|
||||
wait: bool = True,
|
||||
sync_ltc_data: bool = True,
|
||||
) -> None: ...
|
||||
def _get_tensor_id(tensor: Tensor) -> int: ...
|
||||
def _get_tensors_text(tensors: list[Tensor]) -> str: ...
|
||||
def _get_tensors_dot(tensors: list[Tensor]) -> str: ...
|
||||
def _get_tensors_backend(tensors: list[Tensor]) -> str: ...
|
||||
def _get_force_fallback() -> str: ...
|
||||
def _set_force_fallback(newval: str) -> None: ...
|
||||
def _clear_ir_cache() -> None: ...
|
||||
def _dump_ir_cache(filename: str) -> None: ...
|
||||
def _set_reuse_ir(val: bool) -> None: ...
|
||||
def _get_default_device_type() -> str: ...
|
||||
@@ -0,0 +1,12 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# defined in torch/csrc/lazy/python/init.cpp
|
||||
|
||||
from typing import Any
|
||||
|
||||
from torch import Tensor
|
||||
|
||||
def _init(): ...
|
||||
def _get_tensors_ts_device_data_node(
|
||||
tensors: list[Tensor],
|
||||
) -> tuple[list[int], list[Any]]: ...
|
||||
def _run_cached_graph(hash_str: str, graph_inputs: list[Any]) -> list[Tensor]: ...
|
||||
@@ -0,0 +1,58 @@
|
||||
# Defined in torch/csrc/monitor/python_init.cpp
|
||||
|
||||
import datetime
|
||||
from collections.abc import Callable
|
||||
from enum import Enum
|
||||
from types import TracebackType
|
||||
|
||||
class Aggregation(Enum):
|
||||
VALUE = ...
|
||||
MEAN = ...
|
||||
COUNT = ...
|
||||
SUM = ...
|
||||
MAX = ...
|
||||
MIN = ...
|
||||
|
||||
class Stat:
|
||||
name: str
|
||||
count: int
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
aggregations: list[Aggregation],
|
||||
window_size: int,
|
||||
max_samples: int = -1,
|
||||
) -> None: ...
|
||||
def add(self, v: float) -> None: ...
|
||||
def get(self) -> dict[Aggregation, float]: ...
|
||||
|
||||
class Event:
|
||||
name: str
|
||||
timestamp: datetime.datetime
|
||||
data: dict[str, int | float | bool | str]
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
timestamp: datetime.datetime,
|
||||
data: dict[str, int | float | bool | str],
|
||||
) -> None: ...
|
||||
|
||||
def log_event(e: Event) -> None: ...
|
||||
|
||||
class EventHandlerHandle: ...
|
||||
|
||||
def register_event_handler(handler: Callable[[Event], None]) -> EventHandlerHandle: ...
|
||||
def unregister_event_handler(handle: EventHandlerHandle) -> None: ...
|
||||
|
||||
class _WaitCounterTracker:
|
||||
def __enter__(self) -> None: ...
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None = None,
|
||||
exc_value: BaseException | None = None,
|
||||
traceback: TracebackType | None = None,
|
||||
) -> None: ...
|
||||
|
||||
class _WaitCounter:
|
||||
def __init__(self, key: str) -> None: ...
|
||||
def guard(self) -> _WaitCounterTracker: ...
|
||||
@@ -0,0 +1,347 @@
|
||||
# @generated by tools/pyi/gen_pyi.py from torch/_C/_nn.pyi.in
|
||||
# mypy: disable-error-code="type-arg"
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Literal, overload
|
||||
|
||||
from torch import memory_format, Tensor
|
||||
from torch.types import _bool, _device, _dtype, _int, _size
|
||||
|
||||
# Defined in tools/autograd/templates/python_nn_functions.cpp
|
||||
|
||||
def adaptive_avg_pool2d(input: Tensor, output_size: _int | _size) -> Tensor: ...
|
||||
def adaptive_avg_pool3d(input: Tensor, output_size: _int | _size) -> Tensor: ...
|
||||
def adaptive_max_pool2d(
|
||||
input: Tensor,
|
||||
output_size: _int | _size,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def adaptive_max_pool3d(
|
||||
input: Tensor,
|
||||
output_size: _int | _size,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def avg_pool2d(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
stride: _int | _size | None = None,
|
||||
padding: _int | _size = 0,
|
||||
ceil_mode: bool = False,
|
||||
count_include_pad: bool = True,
|
||||
divisor_override: int | None = None,
|
||||
) -> Tensor: ...
|
||||
def avg_pool3d(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
stride: _int | _size | None = None,
|
||||
padding: _int | _size = 0,
|
||||
ceil_mode: bool = False,
|
||||
count_include_pad: bool = True,
|
||||
divisor_override: int | None = None,
|
||||
) -> Tensor: ...
|
||||
def binary_cross_entropy(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
weight: Tensor | None = None,
|
||||
reduction: str = ...,
|
||||
) -> Tensor: ...
|
||||
def col2im(
|
||||
input: Tensor,
|
||||
output_size: _int | _size,
|
||||
kernel_size: _int | _size,
|
||||
dilation: _int | _size,
|
||||
stride: _int | _size | None = None,
|
||||
padding: _int | _size = 0,
|
||||
) -> Tensor: ...
|
||||
def cross_entropy_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
weight: Tensor | None = None,
|
||||
reduction: str = ...,
|
||||
ignore_index: int = -100,
|
||||
label_smoothing: float = 0.0,
|
||||
) -> Tensor: ...
|
||||
def elu(
|
||||
input: Tensor,
|
||||
alpha: float = 1.0,
|
||||
scale: float = 1.0,
|
||||
input_scale: float = 1.0,
|
||||
) -> Tensor: ...
|
||||
def elu_(input: Tensor, alpha: float = ...) -> Tensor: ...
|
||||
def fractional_max_pool2d(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
output_size: _int | _size,
|
||||
_random_samples: Tensor,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def fractional_max_pool3d(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
output_size: _int | _size,
|
||||
_random_samples: Tensor,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def gelu(input: Tensor, approximate: str = ...) -> Tensor: ...
|
||||
def glu(input: Tensor, dim: int = -1) -> Tensor: ...
|
||||
def hardsigmoid(input: Tensor, *, out: Tensor | None = None) -> Tensor: ...
|
||||
def hardsigmoid_(input: Tensor) -> Tensor: ...
|
||||
def hardswish(input: Tensor) -> Tensor: ...
|
||||
def hardswish_(input: Tensor) -> Tensor: ...
|
||||
def hardtanh(
|
||||
input: Tensor,
|
||||
min_val: float = ...,
|
||||
max_val: float = ...,
|
||||
*,
|
||||
out: Tensor | None = None,
|
||||
) -> Tensor: ...
|
||||
def hardtanh_(
|
||||
input: Tensor,
|
||||
min_val: float = ...,
|
||||
max_val: float = ...,
|
||||
) -> Tensor: ...
|
||||
def huber_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
reduction: str = ...,
|
||||
delta: float = 1.0,
|
||||
) -> Tensor: ...
|
||||
def im2col(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
dilation: _int | _size,
|
||||
padding: _int | _size,
|
||||
stride: _int | _size,
|
||||
) -> Tensor: ...
|
||||
def l1_loss(input: Tensor, target: Tensor, reduction: str = ...) -> Tensor: ...
|
||||
def leaky_relu(
|
||||
input: Tensor,
|
||||
negative_slope: float = ...,
|
||||
*,
|
||||
out: Tensor | None = None,
|
||||
) -> Tensor: ...
|
||||
def leaky_relu_(input: Tensor, negative_slope: float = ...) -> Tensor: ...
|
||||
def linear(
|
||||
input: Tensor,
|
||||
weight: Tensor,
|
||||
bias: Tensor | None = None,
|
||||
) -> Tensor: ...
|
||||
def log_sigmoid(input: Tensor) -> Tensor: ...
|
||||
def max_pool2d_with_indices(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
stride: _int | _size | None = None,
|
||||
padding: _int | _size = 0,
|
||||
dilation: _int | _size = 1,
|
||||
ceil_mode: bool = False,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def max_pool3d_with_indices(
|
||||
input: Tensor,
|
||||
kernel_size: _int | _size,
|
||||
stride: _int | _size | None = None,
|
||||
padding: _int | _size = 0,
|
||||
dilation: _int | _size = 1,
|
||||
ceil_mode: bool = False,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
def max_unpool2d(
|
||||
input: Tensor,
|
||||
indices: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
) -> Tensor: ...
|
||||
def max_unpool3d(
|
||||
input: Tensor,
|
||||
indices: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
stride: _int | _size,
|
||||
padding: _int | _size,
|
||||
) -> Tensor: ...
|
||||
def mish(input: Tensor) -> Tensor: ...
|
||||
def mish_(input: Tensor) -> Tensor: ...
|
||||
def mse_loss(input: Tensor, target: Tensor, reduction: str = ...) -> Tensor: ...
|
||||
def multi_margin_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
p: float = 1.0,
|
||||
margin: float = 1.0,
|
||||
weight: Tensor | None = None,
|
||||
reduction: str = ...,
|
||||
) -> Tensor: ...
|
||||
def multilabel_margin_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
reduction: str = ...,
|
||||
) -> Tensor: ...
|
||||
def nll_loss_nd(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
weight: Tensor | None = None,
|
||||
reduction: str = ...,
|
||||
ignore_index: int = -100,
|
||||
) -> Tensor: ...
|
||||
def one_hot(tensor: Tensor, num_classes: int = ...) -> Tensor: ...
|
||||
def pad(
|
||||
input: Tensor,
|
||||
pad: Sequence[int],
|
||||
mode: str = ...,
|
||||
value: float | None = None,
|
||||
) -> Tensor: ...
|
||||
def relu6(input: Tensor) -> Tensor: ...
|
||||
def relu6_(input: Tensor) -> Tensor: ...
|
||||
def scaled_dot_product_attention(
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
value: Tensor,
|
||||
attn_mask: Tensor | None = None,
|
||||
dropout_p: float = 0.0,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
enable_gqa: bool = False,
|
||||
) -> Tensor: ...
|
||||
def silu(input: Tensor) -> Tensor: ...
|
||||
def silu_(input: Tensor) -> Tensor: ...
|
||||
def smooth_l1_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
reduction: str = ...,
|
||||
beta: float = 1.0,
|
||||
) -> Tensor: ...
|
||||
def soft_margin_loss(
|
||||
input: Tensor,
|
||||
target: Tensor,
|
||||
reduction: str = ...,
|
||||
) -> Tensor: ...
|
||||
def softplus(
|
||||
input: Tensor,
|
||||
beta: float = ...,
|
||||
threshold: float = ...,
|
||||
) -> Tensor: ...
|
||||
def softshrink(input: Tensor, lambd: float = ...) -> Tensor: ...
|
||||
|
||||
# Defined in aten/src/ATen/native/mkldnn/Linear.cpp
|
||||
def mkldnn_linear(input: Tensor, weight: Tensor, bias: Tensor | None) -> Tensor: ...
|
||||
|
||||
# Defined at aten/src/ATen/native/mkldnn/MKLDNNConversions.cpp
|
||||
def mkldnn_reorder_conv2d_weight(
|
||||
self: Tensor,
|
||||
padding: list,
|
||||
stride: list,
|
||||
dilatation: list,
|
||||
groups: int,
|
||||
) -> Tensor: ...
|
||||
def mkldnn_reorder_conv3d_weight(
|
||||
self: Tensor,
|
||||
padding: list,
|
||||
stride: list,
|
||||
dilatation: list,
|
||||
groups: int,
|
||||
) -> Tensor: ...
|
||||
|
||||
# Defined in aten/src/ATen/native/mkldnn/Prelu.cpp
|
||||
def mkldnn_prelu(input: Tensor, weight: Tensor) -> Tensor: ...
|
||||
|
||||
# Defined at tools/autograd/templates/python_nn_functions.cpp
|
||||
@overload
|
||||
def _parse_to(
|
||||
device: _device,
|
||||
dtype: _dtype,
|
||||
non_blocking: _bool,
|
||||
copy: _bool,
|
||||
*,
|
||||
memory_format: memory_format,
|
||||
) -> tuple[_device, _dtype, _bool, memory_format]: ...
|
||||
@overload
|
||||
def _parse_to(
|
||||
dtype: _dtype,
|
||||
non_blocking: _bool,
|
||||
copy: _bool,
|
||||
*,
|
||||
memory_format: memory_format,
|
||||
) -> tuple[_device, _dtype, _bool, memory_format]: ...
|
||||
@overload
|
||||
def _parse_to(
|
||||
tensor: Tensor,
|
||||
non_blocking: _bool,
|
||||
copy: _bool,
|
||||
*,
|
||||
memory_format: memory_format,
|
||||
) -> tuple[_device, _dtype, _bool, memory_format]: ...
|
||||
|
||||
# Defined in aten/src/ATen/native/PackedSequence.cpp
|
||||
def pad_sequence(
|
||||
sequences: list[Tensor] | tuple[Tensor, ...],
|
||||
batch_first: bool = False,
|
||||
padding_value: float = 0.0,
|
||||
padding_side: Literal["left", "right"] = "right",
|
||||
) -> Tensor: ...
|
||||
|
||||
# Upsample functions used by torch.nn.functional.interpolate
|
||||
def upsample_nearest1d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_nearest2d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_nearest3d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_nearest_exact1d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_nearest_exact2d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_nearest_exact3d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_linear1d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_bilinear2d_aa(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_bilinear2d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_trilinear3d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_bicubic2d_aa(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def _upsample_lanczos2d_aa(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def upsample_bicubic2d(
|
||||
input: Tensor,
|
||||
output_size: Sequence[int] | None,
|
||||
align_corners: bool,
|
||||
scale_factors: Sequence[float] | None,
|
||||
) -> Tensor: ...
|
||||
def flatten_dense_tensors(tensors: list[Tensor]) -> Tensor: ...
|
||||
def unflatten_dense_tensors(flat: Tensor, tensors: list[Tensor]) -> list[Tensor]: ...
|
||||
@@ -0,0 +1,9 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# Defined in torch/csrc/cuda/shared/nvtx.cpp
|
||||
def rangePushA(message: str) -> int: ...
|
||||
def rangePop() -> int: ...
|
||||
def rangeStartA(message: str) -> int: ...
|
||||
def rangeEnd(int) -> None: ...
|
||||
def markA(message: str) -> None: ...
|
||||
def deviceRangeStart(message: str, stream: int) -> object: ...
|
||||
def deviceRangeEnd(range_handle: object, stream: int) -> None: ...
|
||||
@@ -0,0 +1,39 @@
|
||||
# Defined in torch/csrc/onnx/init.cpp
|
||||
|
||||
from enum import Enum
|
||||
|
||||
PRODUCER_VERSION: str
|
||||
|
||||
class TensorProtoDataType(Enum):
|
||||
UNDEFINED = ...
|
||||
FLOAT = ...
|
||||
UINT8 = ...
|
||||
INT8 = ...
|
||||
UINT16 = ...
|
||||
INT16 = ...
|
||||
INT32 = ...
|
||||
INT64 = ...
|
||||
STRING = ...
|
||||
BOOL = ...
|
||||
FLOAT16 = ...
|
||||
DOUBLE = ...
|
||||
UINT32 = ...
|
||||
UINT64 = ...
|
||||
COMPLEX64 = ...
|
||||
COMPLEX128 = ...
|
||||
BFLOAT16 = ...
|
||||
FLOAT8E5M2 = ...
|
||||
FLOAT8E4M3FN = ...
|
||||
FLOAT8E5M2FNUZ = ...
|
||||
FLOAT8E4M3FNUZ = ...
|
||||
|
||||
class OperatorExportTypes(Enum):
|
||||
ONNX = ...
|
||||
ONNX_ATEN = ...
|
||||
ONNX_ATEN_FALLBACK = ...
|
||||
ONNX_FALLTHROUGH = ...
|
||||
|
||||
class TrainingMode(Enum):
|
||||
EVAL = ...
|
||||
PRESERVE = ...
|
||||
TRAINING = ...
|
||||
@@ -0,0 +1,248 @@
|
||||
from enum import Enum
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
from torch._C import device, dtype, layout
|
||||
|
||||
# defined in torch/csrc/profiler/python/init.cpp
|
||||
|
||||
class RecordScope(Enum):
|
||||
FUNCTION = ...
|
||||
BACKWARD_FUNCTION = ...
|
||||
TORCHSCRIPT_FUNCTION = ...
|
||||
KERNEL_FUNCTION_DTYPE = ...
|
||||
CUSTOM_CLASS = ...
|
||||
BUILD_FEATURE = ...
|
||||
LITE_INTERPRETER = ...
|
||||
USER_SCOPE = ...
|
||||
STATIC_RUNTIME_OP = ...
|
||||
STATIC_RUNTIME_MODEL = ...
|
||||
|
||||
class ProfilerState(Enum):
|
||||
Disabled = ...
|
||||
CPU = ...
|
||||
CUDA = ...
|
||||
NVTX = ...
|
||||
ITT = ...
|
||||
PRIVATEUSE1 = ...
|
||||
KINETO = ...
|
||||
KINETO_GPU_FALLBACK = ...
|
||||
KINETO_PRIVATEUSE1_FALLBACK = ...
|
||||
KINETO_PRIVATEUSE1 = ...
|
||||
|
||||
class ActiveProfilerType(Enum):
|
||||
NONE = ...
|
||||
LEGACY = ...
|
||||
KINETO = ...
|
||||
NVTX = ...
|
||||
ITT = ...
|
||||
PRIVATEUSE1 = ...
|
||||
|
||||
class ProfilerActivity(Enum):
|
||||
CPU = ...
|
||||
CUDA = ...
|
||||
XPU = ...
|
||||
MTIA = ...
|
||||
HPU = ...
|
||||
PrivateUse1 = ...
|
||||
|
||||
class _EventType(Enum):
|
||||
TorchOp = ...
|
||||
Backend = ...
|
||||
Allocation = ...
|
||||
OutOfMemory = ...
|
||||
PyCall = ...
|
||||
PyCCall = ...
|
||||
Kineto = ...
|
||||
|
||||
class _ExperimentalConfig:
|
||||
def __init__(
|
||||
self,
|
||||
profiler_metrics: list[str] = ...,
|
||||
profiler_measure_per_kernel: bool = ...,
|
||||
verbose: bool = ...,
|
||||
performance_events: list[str] = ...,
|
||||
enable_cuda_sync_events: bool = ...,
|
||||
profile_all_threads: bool = ...,
|
||||
) -> None: ...
|
||||
|
||||
class ProfilerConfig:
|
||||
def __init__(
|
||||
self,
|
||||
state: ProfilerState,
|
||||
report_input_shapes: bool,
|
||||
profile_memory: bool,
|
||||
with_stack: bool,
|
||||
with_flops: bool,
|
||||
with_modules: bool,
|
||||
experimental_config: _ExperimentalConfig,
|
||||
trace_id: str | None = None,
|
||||
) -> None: ...
|
||||
|
||||
class _ProfilerEvent:
|
||||
start_tid: int
|
||||
start_time_ns: int
|
||||
children: list[_ProfilerEvent]
|
||||
|
||||
# TODO(robieta): remove in favor of `self.typed`
|
||||
extra_fields: (
|
||||
_ExtraFields_TorchOp
|
||||
| _ExtraFields_Backend
|
||||
| _ExtraFields_Allocation
|
||||
| _ExtraFields_OutOfMemory
|
||||
| _ExtraFields_PyCall
|
||||
| _ExtraFields_PyCCall
|
||||
| _ExtraFields_Kineto
|
||||
)
|
||||
|
||||
@property
|
||||
def typed(
|
||||
self,
|
||||
) -> (
|
||||
tuple[Literal[_EventType.TorchOp], _ExtraFields_TorchOp]
|
||||
| tuple[Literal[_EventType.Backend], _ExtraFields_Backend]
|
||||
| tuple[Literal[_EventType.Allocation], _ExtraFields_Allocation]
|
||||
| tuple[Literal[_EventType.OutOfMemory], _ExtraFields_OutOfMemory]
|
||||
| tuple[Literal[_EventType.PyCall], _ExtraFields_PyCall]
|
||||
| tuple[Literal[_EventType.PyCCall], _ExtraFields_PyCCall]
|
||||
| tuple[Literal[_EventType.Kineto], _ExtraFields_Kineto]
|
||||
): ...
|
||||
@property
|
||||
def name(self) -> str: ...
|
||||
@property
|
||||
def tag(self) -> _EventType: ...
|
||||
@property
|
||||
def id(self) -> int: ...
|
||||
@property
|
||||
def parent(self) -> _ProfilerEvent | None: ...
|
||||
@property
|
||||
def correlation_id(self) -> int: ...
|
||||
@property
|
||||
def end_time_ns(self) -> int: ...
|
||||
@property
|
||||
def duration_time_ns(self) -> int: ...
|
||||
|
||||
class _TensorMetadata:
|
||||
impl_ptr: int | None
|
||||
storage_data_ptr: int | None
|
||||
id: int | None
|
||||
|
||||
@property
|
||||
def allocation_id(self) -> int | None: ...
|
||||
@property
|
||||
def layout(self) -> layout: ...
|
||||
@property
|
||||
def device(self) -> device: ...
|
||||
@property
|
||||
def dtype(self) -> dtype: ...
|
||||
@property
|
||||
def sizes(self) -> list[int]: ...
|
||||
@property
|
||||
def strides(self) -> list[int]: ...
|
||||
|
||||
Scalar: TypeAlias = int | float | bool | complex
|
||||
Input: TypeAlias = _TensorMetadata | list[_TensorMetadata] | Scalar | None
|
||||
|
||||
class _ExtraFields_TorchOp:
|
||||
name: str
|
||||
sequence_number: int
|
||||
allow_tf32_cublas: bool
|
||||
|
||||
@property
|
||||
def inputs(self) -> list[Input]: ...
|
||||
@property
|
||||
def scope(self) -> RecordScope: ...
|
||||
|
||||
class _ExtraFields_Backend: ...
|
||||
|
||||
class _ExtraFields_Allocation:
|
||||
ptr: int
|
||||
id: int | None
|
||||
alloc_size: int
|
||||
total_allocated: int
|
||||
total_reserved: int
|
||||
|
||||
@property
|
||||
def allocation_id(self) -> int | None: ...
|
||||
@property
|
||||
def device(self) -> device: ...
|
||||
|
||||
class _ExtraFields_OutOfMemory: ...
|
||||
|
||||
class _PyFrameState:
|
||||
line_number: int
|
||||
function_name: str
|
||||
|
||||
@property
|
||||
def file_name(self) -> str: ...
|
||||
|
||||
class _NNModuleInfo:
|
||||
@property
|
||||
def self_ptr(self) -> int: ...
|
||||
@property
|
||||
def cls_ptr(self) -> int: ...
|
||||
@property
|
||||
def cls_name(self) -> str: ...
|
||||
@property
|
||||
def parameters(
|
||||
self,
|
||||
) -> list[tuple[str, _TensorMetadata, _TensorMetadata | None]]: ...
|
||||
|
||||
class _OptimizerInfo:
|
||||
@property
|
||||
def parameters(
|
||||
self,
|
||||
) -> list[
|
||||
tuple[
|
||||
# Parameter
|
||||
_TensorMetadata,
|
||||
#
|
||||
# Gradient (if present during optimizer.step())
|
||||
_TensorMetadata | None,
|
||||
#
|
||||
# Optimizer state for Parameter as (name, tensor) pairs
|
||||
list[tuple[str, _TensorMetadata]],
|
||||
]
|
||||
]: ...
|
||||
|
||||
class _ExtraFields_PyCCall:
|
||||
@property
|
||||
def caller(self) -> _PyFrameState: ...
|
||||
|
||||
class _ExtraFields_PyCall:
|
||||
@property
|
||||
def callsite(self) -> _PyFrameState: ...
|
||||
@property
|
||||
def caller(self) -> _PyFrameState: ...
|
||||
@property
|
||||
def module(self) -> _NNModuleInfo | None: ...
|
||||
@property
|
||||
def optimizer(self) -> _OptimizerInfo | None: ...
|
||||
|
||||
class _ExtraFields_Kineto: ...
|
||||
|
||||
def _add_execution_trace_observer(output_file_path: str) -> bool: ...
|
||||
def _remove_execution_trace_observer() -> None: ...
|
||||
def _enable_execution_trace_observer() -> None: ...
|
||||
def _disable_execution_trace_observer() -> None: ...
|
||||
def _set_record_concrete_inputs_enabled_val(val: bool) -> None: ...
|
||||
def _set_fwd_bwd_enabled_val(val: bool) -> None: ...
|
||||
def _set_cuda_sync_enabled_val(val: bool) -> None: ...
|
||||
|
||||
class CapturedTraceback: ...
|
||||
|
||||
def gather_traceback(python: bool, script: bool, cpp: bool) -> CapturedTraceback: ...
|
||||
|
||||
# The Dict has name, filename, line
|
||||
def symbolize_tracebacks(
|
||||
to_symbolize: list[CapturedTraceback],
|
||||
) -> list[list[dict[str, str]]]: ...
|
||||
|
||||
class _RecordFunctionFast:
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
input_values: list | tuple | None = None,
|
||||
keyword_values: dict | None = None,
|
||||
) -> None: ...
|
||||
def __enter__(self) -> None: ...
|
||||
def __exit__(self, *exc_info: object) -> None: ...
|
||||
@@ -0,0 +1,3 @@
|
||||
# Defined in torch/csrc/utils/verbose.cpp
|
||||
def mkl_set_verbose(enable: int) -> int: ...
|
||||
def mkldnn_set_verbose(level: int) -> int: ...
|
||||
Reference in New Issue
Block a user