xref: /aosp_15_r20/external/pytorch/torch/types.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# mypy: allow-untyped-defs
2
3# In some cases, these basic types are shadowed by corresponding
4# top-level values.  The underscore variants let us refer to these
5# types.  See https://github.com/python/mypy/issues/4146 for why these
6# workarounds is necessary
7from builtins import (  # noqa: F401
8    bool as _bool,
9    bytes as _bytes,
10    complex as _complex,
11    float as _float,
12    int as _int,
13    str as _str,
14)
15from typing import Any, Dict, List, Sequence, Tuple, TYPE_CHECKING, Union
16from typing_extensions import TypeAlias
17
18# `as` imports have better static analysis support than assignment `ExposedType: TypeAlias = HiddenType`
19from torch import (  # noqa: F401
20    device as _device,
21    DispatchKey as DispatchKey,
22    dtype as _dtype,
23    layout as _layout,
24    qscheme as _qscheme,
25    Size as Size,
26    SymBool as SymBool,
27    SymFloat as SymFloat,
28    SymInt as SymInt,
29    Tensor as Tensor,
30)
31
32
33if TYPE_CHECKING:
34    from torch.autograd.graph import GradientEdge
35
36
37__all__ = ["Number", "Device", "Storage"]
38
39# Convenience aliases for common composite types that we need
40# to talk about in PyTorch
41_TensorOrTensors: TypeAlias = Union[Tensor, Sequence[Tensor]]  # noqa: PYI047
42_TensorOrTensorsOrGradEdge: TypeAlias = Union[  # noqa: PYI047
43    Tensor,
44    Sequence[Tensor],
45    "GradientEdge",
46    Sequence["GradientEdge"],
47]
48
49_size: TypeAlias = Union[Size, List[int], Tuple[int, ...]]  # noqa: PYI042,PYI047
50_symsize: TypeAlias = Union[Size, Sequence[Union[int, SymInt]]]  # noqa: PYI042,PYI047
51_dispatchkey: TypeAlias = Union[str, DispatchKey]  # noqa: PYI042,PYI047
52
53# int or SymInt
54IntLikeType: TypeAlias = Union[int, SymInt]
55# float or SymFloat
56FloatLikeType: TypeAlias = Union[float, SymFloat]
57# bool or SymBool
58BoolLikeType: TypeAlias = Union[bool, SymBool]
59
60py_sym_types = (SymInt, SymFloat, SymBool)
61PySymType: TypeAlias = Union[SymInt, SymFloat, SymBool]
62
63# Meta-type for "numeric" things; matches our docs
64Number: TypeAlias = Union[int, float, bool]
65
66# Meta-type for "device-like" things.  Not to be confused with 'device' (a
67# literal device object).  This nomenclature is consistent with PythonArgParser.
68# None means use the default device (typically CPU)
69Device: TypeAlias = Union[_device, str, int, None]
70
71
72# Storage protocol implemented by ${Type}StorageBase classes
73class Storage:
74    _cdata: int
75    device: _device
76    dtype: _dtype
77    _torch_load_uninitialized: bool
78
79    def __deepcopy__(self, memo: Dict[int, Any]) -> "Storage":
80        raise NotImplementedError
81
82    def _new_shared(self, size: int) -> "Storage":
83        raise NotImplementedError
84
85    def _write_file(
86        self,
87        f: Any,
88        is_real_file: bool,
89        save_size: bool,
90        element_size: int,
91    ) -> None:
92        raise NotImplementedError
93
94    def element_size(self) -> int:
95        raise NotImplementedError
96
97    def is_shared(self) -> bool:
98        raise NotImplementedError
99
100    def share_memory_(self) -> "Storage":
101        raise NotImplementedError
102
103    def nbytes(self) -> int:
104        raise NotImplementedError
105
106    def cpu(self) -> "Storage":
107        raise NotImplementedError
108
109    def data_ptr(self) -> int:
110        raise NotImplementedError
111
112    def from_file(
113        self,
114        filename: str,
115        shared: bool = False,
116        nbytes: int = 0,
117    ) -> "Storage":
118        raise NotImplementedError
119
120    def _new_with_file(
121        self,
122        f: Any,
123        element_size: int,
124    ) -> "Storage":
125        raise NotImplementedError
126