pytorch

1
# mypy: allow-untyped-defs
2
from typing import Any
3

4
import torch
5

6
class GlobalStateGuard:
7
    def check(self) -> bool: ...
8
    def reason(self) -> str: ...
9

10
class LeafGuard: ...
11
class GuardDebugInfo: ...
12

13
class GuardManager:
14
    def check(self, value) -> bool: ...
15
    def check_verbose(self, value) -> GuardDebugInfo: ...
16

17
    # Accessors
18
    def globals_dict_manager(
19
        self,
20
        f_globals: dict[str, Any],
21
        source,
22
        example_value,
23
        guard_manager_enum,
24
    ) -> GuardManager: ...
25
    def dict_getitem_manager(
26
        self,
27
        key,
28
        source,
29
        example_value,
30
        guard_manager_enum,
31
    ) -> GuardManager: ...
32
    def global_weakref_manager(
33
        self,
34
        global_name: str,
35
        source,
36
        example_value,
37
        guard_manager_enum,
38
    ) -> GuardManager: ...
39
    def type_manager(
40
        self,
41
        source,
42
        example_value,
43
        guard_manager_enum,
44
    ) -> GuardManager: ...
45
    def getattr_manager(
46
        self,
47
        attr: str,
48
        source,
49
        example_value,
50
        guard_manager_enum,
51
    ) -> GuardManager: ...
52
    def lambda_manager(
53
        self,
54
        python_lambda,
55
        source,
56
        example_value,
57
        guard_manager_enum,
58
    ) -> GuardManager: ...
59

60
    # Leaf guards
61
    def add_lambda_guard(self, user_lambda, verbose_code_parts: list[str]) -> None: ...
62
    def add_id_match_guard(self, id_val, verbose_code_parts: list[str]) -> None: ...
63
    def add_equals_match_guard(
64
        self,
65
        equals_val,
66
        verbose_code_parts: list[str],
67
    ) -> None: ...
68
    def add_global_state_guard(self, verbose_code_parts: list[str]) -> None: ...
69
    def add_torch_function_mode_stack_guard(
70
        self, initial_stack, ignored_types, verbose_code_parts: list[str]
71
    ) -> None: ...
72

73
class RootGuardManager(GuardManager):
74
    def get_epilogue_lambda_guards(self) -> list[LeafGuard]: ...
75
    def add_epilogue_lambda_guard(
76
        self,
77
        guard: LeafGuard,
78
        verbose_code_parts: list[str],
79
    ) -> None: ...
80

81
class DictGuardManager(GuardManager):
82
    def get_key_manager(
83
        self,
84
        index,
85
        source,
86
        example_value,
87
        guard_manager_enum,
88
    ) -> GuardManager: ...
89
    def get_value_manager(
90
        self,
91
        index,
92
        source,
93
        example_value,
94
        guard_manager_enum,
95
    ) -> GuardManager: ...
96

97
def install_object_aliasing_guard(
98
    guard_managers: list[GuardManager],
99
    tensor_names: list[str],
100
    verbose_code_parts: list[str],
101
): ...
102
def install_no_tensor_aliasing_guard(
103
    guard_managers: list[GuardManager],
104
    tensor_names: list[str],
105
    verbose_code_parts: list[str],
106
): ...
107

108
class TensorGuards:
109
    def __init__(
110
        self,
111
        *,
112
        dynamic_dims_sizes: list[torch.SymInt | None] | None = None,
113
        dynamic_dims_strides: list[torch.SymInt | None] | None = None,
114
    ) -> None: ...
115
    def check(self, *args) -> bool: ...
116
    def check_verbose(self, *args, tensor_check_names=None) -> bool | str: ...
117

118
def assert_size_stride(
119
    item: torch.Tensor,
120
    size: torch.types._size,
121
    stride: torch.types._size,
122
): ...
123
def check_obj_id(obj: object, expected: int) -> bool: ...
124
def check_type_id(obj: object, expected: int) -> bool: ...
125
def dict_version(d: dict[Any, Any]) -> int: ...
126