functional.pyi 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434
  1. from torch import Tensor
  2. from torch.types import _size, _dtype
  3. from typing import Any, Optional, Tuple, Dict, List, Callable, Sequence, Union
  4. from .common_types import _ratio_any_t, _size_any_t, _size_1_t, _size_2_t, _size_3_t, _size_2_opt_t, _size_3_opt_t
  5. # 'TypedDict' is a new accepted type that represents a dictionary with a fixed set of allowed keys.
  6. # It is standards-track but not in `typing` yet. We leave this hear to be uncommented once the feature
  7. # is wide-spread.
  8. # from mypy_extensions import TypedDict
  9. # GRID_SAMPLE_INTERPOLATION_MODES = TypedDict('GRID_SAMPLE_INTERPOLATION_MODES', {'bilinear': int, 'nearest': int})
  10. # GRID_SAMPLE_PADDING_MODES = TypedDict('GRID_SAMPLE_PADDING_MODES', {'zeros': int, 'border': int, 'reflection': int})
  11. GRID_SAMPLE_INTERPOLATION_MODES = Dict[str, int]
  12. GRID_SAMPLE_PADDING_MODES = Dict[str, int]
  13. # These stubs were generated by running stubgen (`stubgen --parse-only functional.py`), followed by manual cleaning.
  14. #
  15. # The 'BroadcastingList{1,2,3}' types were replaced by `_size` or _output_ratio, as appropriate.
  16. # This was necessary since the JIT uses BroadcastingList* types but static checking with mypy etc requires a `Sequence`
  17. # type. There is no way to express the expected lengths of these lists in the current Python typing system.
  18. #
  19. # Functions created via `_add_docstr` in `functional.py` where merely typed as `Any` by `stubgen`, so those were
  20. # deleted from the stub and replaced by generated declarations. See `gen_pyi` for the implementation of the code
  21. # generation logic for those functions. In the future, it might be worth looking into using the mypy plugin system
  22. # to encode the type semantics of `_add_docstr`, should that system ever become widespread.
  23. def fractional_max_pool2d_with_indices(input: Tensor, kernel_size: _size, output_size: Optional[_size] = ...,
  24. output_ratio: Optional[_ratio_any_t] = ..., return_indices: bool = ...,
  25. _random_samples: Optional[Tensor] = ...) -> Tuple[Tensor, Tensor]: ...
  26. def fractional_max_pool3d_with_indices(input: Tensor, kernel_size: _size, output_size: Optional[_size] = ...,
  27. output_ratio: Optional[_ratio_any_t] = ..., return_indices: bool = ...,
  28. _random_samples: Optional[Tensor] = ...) -> Tuple[Tensor, Tensor]: ...
  29. def max_pool1d_with_indices(input: Tensor, kernel_size: _size, stride: Optional[_size] = ..., padding: _size = ...,
  30. dilation: _size = ..., ceil_mode: bool = ..., return_indices: bool = ...) -> Tuple[
  31. Tensor, Tensor]: ...
  32. def max_pool2d_with_indices(input: Tensor, kernel_size: _size, stride: Optional[_size] = ..., padding: _size = ...,
  33. dilation: _size = ..., ceil_mode: bool = ..., return_indices: bool = ...) -> Tuple[
  34. Tensor, Tensor]: ...
  35. def max_pool3d_with_indices(input: Tensor, kernel_size: _size, stride: Optional[_size] = ..., padding: _size = ...,
  36. dilation: _size = ..., ceil_mode: bool = ..., return_indices: bool = ...) -> Tuple[
  37. Tensor, Tensor]: ...
  38. def max_unpool1d(input: Tensor, indices: Tensor, kernel_size: _size, stride: Optional[_size] = ...,
  39. padding: _size = ..., output_size: Optional[_size] = ...) -> Tensor: ...
  40. def max_unpool2d(input: Tensor, indices: Tensor, kernel_size: _size, stride: Optional[_size] = ...,
  41. padding: _size = ..., output_size: Optional[_size] = ...) -> Tensor: ...
  42. def max_unpool3d(input: Tensor, indices: Tensor, kernel_size: _size, stride: Optional[_size] = ...,
  43. padding: _size = ..., output_size: Optional[_size] = ...) -> Tensor: ...
  44. def lp_pool1d(input: Tensor, norm_type: float, kernel_size: _size_1_t, stride: Union[Optional[_size], Optional[int]] = ...,
  45. ceil_mode: bool = ...) -> Tensor: ...
  46. def lp_pool2d(input: Tensor, norm_type: float, kernel_size: _size_2_t, stride: Union[Optional[_size], Optional[int]] = ...,
  47. ceil_mode: bool = ...) -> Tensor: ...
  48. def adaptive_max_pool1d_with_indices(input: Tensor, output_size: _size, return_indices: bool = ...) -> Tuple[
  49. Tensor, Tensor]: ...
  50. def adaptive_max_pool2d_with_indices(input: Tensor, output_size: _size_2_opt_t, return_indices: bool = ...) -> Tuple[
  51. Tensor, Tensor]: ...
  52. def adaptive_max_pool3d_with_indices(input: Tensor, output_size: _size_3_opt_t, return_indices: bool = ...) -> Tuple[
  53. Tensor, Tensor]: ...
  54. def adaptive_avg_pool1d(input: Tensor, output_size: _size_1_t) -> Tensor: ...
  55. def adaptive_avg_pool2d(input: Tensor, output_size: _size_2_opt_t) -> Tensor: ...
  56. def adaptive_avg_pool3d(input: Tensor, output_size: _size_3_opt_t) -> Tensor: ...
  57. def dropout(input: Tensor, p: float = ..., training: bool = ..., inplace: bool = ...) -> Tensor: ...
  58. def alpha_dropout(input: Tensor, p: float = ..., training: bool = ..., inplace: bool = ...) -> Tensor: ...
  59. def dropout1d(input: Tensor, p: float = ..., training: bool = ..., inplace: bool = ...) -> Tensor: ...
  60. def dropout2d(input: Tensor, p: float = ..., training: bool = ..., inplace: bool = ...) -> Tensor: ...
  61. def dropout3d(input: Tensor, p: float = ..., training: bool = ..., inplace: bool = ...) -> Tensor: ...
  62. def feature_alpha_dropout(input: Tensor, p: float = ..., training: bool = ..., inplace: bool = ...) -> Tensor: ...
  63. def threshold(input: Tensor, threshold: float, value: float, inplace: bool = ...) -> Tensor: ...
  64. def relu(input: Tensor, inplace: bool = ...) -> Tensor: ...
  65. def glu(input: Tensor, dim: int = ...) -> Tensor: ...
  66. def hardtanh(input: Tensor, min_val: float = ..., max_val: float = ..., inplace: bool = ...) -> Tensor: ...
  67. def relu6(input: Tensor, inplace: bool = ...) -> Tensor: ...
  68. def elu(input: Tensor, alpha: float = ..., inplace: bool = ...) -> Tensor: ...
  69. def selu(input: Tensor, inplace: bool = ...) -> Tensor: ...
  70. def celu(input: Tensor, alpha: float = ..., inplace: bool = ...) -> Tensor: ...
  71. def leaky_relu(input: Tensor, negative_slope: float = ..., inplace: bool = ...) -> Tensor: ...
  72. def prelu(input: Tensor, weight: Tensor) -> Tensor: ...
  73. def rrelu(input: Tensor, lower: float = ..., upper: float = ..., training: bool = ...,
  74. inplace: bool = ...) -> Tensor: ...
  75. def gelu(input: Any, approximate: str = ...): ...
  76. def hardshrink(input: Tensor, lambd: float = ...) -> Tensor: ...
  77. def tanhshrink(input: Any): ...
  78. def softsign(input: Any): ...
  79. def softmin(input: Tensor, dim: Optional[int] = ..., _stacklevel: int = ..., dtype: Optional[_dtype] = ...) -> Tensor: ...
  80. def softmax(input: Tensor, dim: Optional[int] = ..., _stacklevel: int = ..., dtype: Optional[_dtype] = ...) -> Tensor: ...
  81. def gumbel_softmax(logits: Tensor, tau: float = ..., hard: bool = ..., eps: float = ..., dim: int = ...) -> Tensor: ...
  82. def log_softmax(input: Tensor, dim: Optional[int] = ..., _stacklevel: int = ...,
  83. dtype: Optional[_dtype] = ...) -> Tensor: ...
  84. def tanh(input: Any): ...
  85. def sigmoid(input: Any) -> Tensor: ...
  86. def hardsigmoid(input: Tensor, inplace: bool = False) -> Tensor: ...
  87. def linear(input: Tensor, weight: Tensor, bias: Optional[Tensor] = ...) -> Tensor: ...
  88. def bilinear(input1: Tensor, input2: Tensor, weight: Tensor, bias: Optional[Tensor] = ...) -> Tensor: ...
  89. def silu(input: Tensor, inplace: bool = False) -> Tensor: ...
  90. def mish(input: Tensor, inplace: bool = False) -> Tensor: ...
  91. def hardswish(input: Tensor, inplace: bool = False) -> Tensor: ...
  92. def embedding(input: Tensor, weight: Tensor, padding_idx: Optional[int] = ..., max_norm: Optional[float] = ...,
  93. norm_type: float = ..., scale_grad_by_freq: bool = ..., sparse: bool = ...) -> Tensor: ...
  94. def embedding_bag(input: Tensor, weight: Tensor, offsets: Optional[Tensor] = ..., max_norm: Optional[float] = ...,
  95. norm_type: float = ..., scale_grad_by_freq: bool = ..., mode: str = ...,
  96. sparse: bool = ..., per_sample_weights: Optional[Tensor] = ...,
  97. include_last_offset: bool = ..., padding_idx: Optional[int] = ...) -> Tensor: ...
  98. def batch_norm(input: Tensor, running_mean: Optional[Tensor], running_var: Optional[Tensor],
  99. weight: Optional[Tensor] = ..., bias: Optional[Tensor] = ..., training: bool = ...,
  100. momentum: float = ..., eps: float = ...) -> Tensor: ...
  101. def instance_norm(input: Tensor, running_mean: Optional[Tensor] = ..., running_var: Optional[Tensor] = ...,
  102. weight: Optional[Tensor] = ..., bias: Optional[Tensor] = ..., use_input_stats: bool = ...,
  103. momentum: float = ..., eps: float = ...) -> Tensor: ...
  104. def layer_norm(input: Tensor, normalized_shape: Sequence[int], weight: Optional[Tensor] = ..., bias: Optional[Tensor] = ...,
  105. eps: float = ...) -> Tensor: ...
  106. def group_norm(input: Tensor, num_groups: int, weight: Optional[Tensor] = ..., bias: Optional[Tensor] = ...,
  107. eps: float = ...) -> Tensor: ...
  108. def local_response_norm(input: Tensor, size: int, alpha: float = ..., beta: float = ..., k: float = ...) -> Tensor: ...
  109. def ctc_loss(log_probs: Tensor, targets: Tensor, input_lengths: Tensor, target_lengths: Tensor, blank: int = ...,
  110. reduction: str = ..., zero_infinity: bool = ...) -> Tensor: ...
  111. def nll_loss(input: Tensor, target: Tensor, weight: Optional[Tensor] = ..., size_average: Optional[bool] = ...,
  112. ignore_index: int = ..., reduce: Optional[bool] = ..., reduction: str = ...) -> Tensor: ...
  113. def poisson_nll_loss(input: Tensor, target: Tensor, log_input: bool = ..., full: bool = ...,
  114. size_average: Optional[bool] = ..., eps: float = ..., reduce: Optional[bool] = ...,
  115. reduction: str = ...) -> Tensor: ...
  116. def gaussian_nll_loss(input: Tensor, target: Tensor, var: Tensor, full: Optional[bool] = ...,
  117. eps: Optional[float] = ..., reduction: Optional[str] = ...) -> Tensor: ...
  118. def kl_div(input: Tensor, target: Tensor, size_average: Optional[bool] = ..., reduce: Optional[bool] = ...,
  119. reduction: str = ..., log_target: bool = ...) -> Tensor: ...
  120. def cross_entropy(input: Tensor, target: Tensor, weight: Optional[Tensor] = ..., size_average: Optional[bool] = ...,
  121. ignore_index: int = ..., reduce: Optional[bool] = ..., reduction: str = ...,
  122. label_smoothing: float = ...) -> Tensor: ...
  123. def binary_cross_entropy(input: Tensor, target: Tensor, weight: Optional[Tensor] = ...,
  124. size_average: Optional[bool] = ..., reduce: Optional[bool] = ...,
  125. reduction: str = ...) -> Tensor: ...
  126. def binary_cross_entropy_with_logits(input: Tensor, target: Tensor, weight: Optional[Tensor] = ...,
  127. size_average: Optional[bool] = ..., reduce: Optional[bool] = ...,
  128. reduction: str = ..., pos_weight: Optional[Tensor] = ...) -> Tensor: ...
  129. def smooth_l1_loss(input: Tensor, target: Tensor, size_average: Optional[bool] = ..., reduce: Optional[bool] = ...,
  130. reduction: str = ..., beta: float = ...) -> Tensor: ...
  131. def huber_loss(input: Tensor, target: Tensor, reduction: str = ..., delta: float = ...) -> Tensor: ...
  132. def l1_loss(input: Tensor, target: Tensor, size_average: Optional[bool] = ..., reduce: Optional[bool] = ...,
  133. reduction: str = ...) -> Tensor: ...
  134. def mse_loss(input: Tensor, target: Tensor, size_average: Optional[bool] = ..., reduce: Optional[bool] = ...,
  135. reduction: str = ...) -> Tensor: ...
  136. def margin_ranking_loss(input1: Tensor, input2: Tensor, target: Tensor, margin: float = ...,
  137. size_average: Optional[bool] = ..., reduce: Optional[bool] = ...,
  138. reduction: str = ...) -> Tensor: ...
  139. def hinge_embedding_loss(input: Tensor, target: Tensor, margin: float = ..., size_average: Optional[bool] = ...,
  140. reduce: Optional[bool] = ..., reduction: str = ...) -> Tensor: ...
  141. def multilabel_margin_loss(input: Tensor, target: Tensor, size_average: Optional[bool] = ...,
  142. reduce: Optional[bool] = ..., reduction: str = ...) -> Tensor: ...
  143. def soft_margin_loss(input: Tensor, target: Tensor, size_average: Optional[bool] = ..., reduce: Optional[bool] = ...,
  144. reduction: str = ...) -> Tensor: ...
  145. def multilabel_soft_margin_loss(input: Tensor, target: Tensor, weight: Optional[Tensor] = ...,
  146. size_average: Optional[bool] = ..., reduce: Optional[bool] = ...,
  147. reduction: str = ...) -> Tensor: ...
  148. def cosine_embedding_loss(input1: Tensor, input2: Tensor, target: Tensor, margin: float = ...,
  149. size_average: Optional[bool] = ..., reduce: Optional[bool] = ...,
  150. reduction: str = ...) -> Tensor: ...
  151. def multi_margin_loss(input: Tensor, target: Tensor, p: int = ..., margin: float = ..., weight: Optional[Tensor] = ...,
  152. size_average: Optional[bool] = ..., reduce: Optional[bool] = ...,
  153. reduction: str = ...) -> Tensor: ...
  154. def upsample(input: Any, size: Optional[Any] = ..., scale_factor: Optional[Any] = ..., mode: str = ...,
  155. align_corners: Optional[Any] = ...): ...
  156. def interpolate(input: Any, size: Optional[Any] = ..., scale_factor: Optional[Any] = ..., mode: str = ...,
  157. align_corners: Optional[Any] = ..., recompute_scale_factor: Optional[Any] = ...,
  158. antialias: bool = ...): ...
  159. def upsample_nearest(input: Any, size: Optional[Any] = ..., scale_factor: Optional[Any] = ...): ...
  160. def upsample_bilinear(input: Any, size: Optional[Any] = ..., scale_factor: Optional[Any] = ...): ...
  161. def grid_sample(input: Tensor, grid: Tensor, mode: str = ..., padding_mode: str = ...,
  162. align_corners: Optional[Any] = ...) -> Tensor: ...
  163. def affine_grid(theta: Tensor, size: List[int], align_corners: Optional[Any] = ...) -> Tensor: ...
  164. def pad(input: Tensor, pad: Sequence[int], mode: str = ..., value: float = ...) -> Tensor: ...
  165. def pairwise_distance(x1: Tensor, x2: Tensor, p: float = ..., eps: float = ..., keepdim: bool = ...) -> Tensor: ...
  166. def triplet_margin_loss(anchor: Tensor, positive: Tensor, negative: Tensor, margin: float = ..., p: float = ...,
  167. eps: float = ..., swap: bool = ..., size_average: Optional[bool] = ...,
  168. reduce: Optional[bool] = ..., reduction: str = ...) -> Tensor: ...
  169. def triplet_margin_with_distance_loss(anchor: Tensor, positive: Tensor, negative: Tensor, *,
  170. distance_function: Optional[Callable[[Tensor, Tensor], Tensor]]=...,
  171. margin: float=..., swap: bool=..., reduction: str=...) -> Tensor: ...
  172. def normalize(input: Tensor, p: float = ..., dim: int = ..., eps: float = ...,
  173. out: Optional[Tensor] = ...) -> Tensor: ...
  174. def assert_int_or_pair(arg: Any, arg_name: Any, message: Any) -> None: ...
  175. def unfold(input: Tensor, kernel_size: _size_any_t, dilation: _size_any_t = ..., padding: _size_any_t = ...,
  176. stride: _size_any_t = ...) -> Tensor: ...
  177. def fold(input: Tensor, output_size: _size_any_t, kernel_size: _size_any_t, dilation: _size_any_t = ..., padding: _size_any_t = ...,
  178. stride: _size_any_t = ...) -> Tensor: ...
  179. def _canonical_mask(
  180. mask: Optional[Tensor],
  181. mask_name: str,
  182. other_type: Optional[_dtype],
  183. other_name: str,
  184. target_type: _dtype,
  185. check_other: bool = True,
  186. ) -> Optional[Tensor]: ...
  187. def _none_or_dtype(input: Optional[Tensor]) -> Optional[_dtype]: ...
  188. def multi_head_attention_forward(query: Tensor,
  189. key: Tensor,
  190. value: Tensor,
  191. embed_dim_to_check: int,
  192. num_heads: int,
  193. in_proj_weight: Optional[Tensor],
  194. in_proj_bias: Optional[Tensor],
  195. bias_k: Optional[Tensor],
  196. bias_v: Optional[Tensor],
  197. add_zero_attn: bool,
  198. dropout_p: float,
  199. out_proj_weight: Tensor,
  200. out_proj_bias: Optional[Tensor],
  201. training: bool = True,
  202. key_padding_mask: Optional[Tensor] = None,
  203. need_weights: bool = True,
  204. attn_mask: Optional[Tensor] = None,
  205. use_separate_proj_weight: bool = False,
  206. q_proj_weight: Optional[Tensor] = None,
  207. k_proj_weight: Optional[Tensor] = None,
  208. v_proj_weight: Optional[Tensor] = None,
  209. static_k: Optional[Tensor] = None,
  210. static_v: Optional[Tensor] = None,
  211. average_attn_weights: bool = True,
  212. is_causal: bool = False
  213. ) -> Tuple[Tensor, Optional[Tensor]]: ...
  214. from .. import conv1d as conv1d
  215. from .. import conv2d as conv2d
  216. from .. import conv3d as conv3d
  217. from .. import conv_transpose1d as conv_transpose1d
  218. from .. import conv_transpose2d as conv_transpose2d
  219. from .. import conv_transpose3d as conv_transpose3d
  220. from .. import conv_tbc as conv_tbc
  221. from .. import avg_pool1d as avg_pool1d
  222. from .. import relu_ as relu_
  223. from .. import selu_ as selu_
  224. from .. import celu_ as celu_
  225. from .. import rrelu_ as rrelu_
  226. from .. import pixel_shuffle as pixel_shuffle
  227. from .. import pixel_unshuffle as pixel_unshuffle
  228. from .. import channel_shuffle as channel_shuffle
  229. from .. import native_channel_shuffle as native_channel_shuffle
  230. from .. import pdist as pdist
  231. from .. import cosine_similarity as cosine_similarity
  232. fractional_max_pool2d: Callable
  233. fractional_max_pool3d: Callable
  234. max_pool1d: Callable
  235. max_pool2d: Callable
  236. max_pool3d: Callable
  237. adaptive_max_pool1d: Callable
  238. adaptive_max_pool2d: Callable
  239. adaptive_max_pool3d: Callable
  240. avg_pool2d: Callable
  241. avg_pool3d: Callable
  242. hardtanh_: Callable
  243. elu_: Callable
  244. leaky_relu_: Callable
  245. logsigmoid: Callable
  246. softplus: Callable
  247. softshrink: Callable
  248. one_hot: Callable
  249. scaled_dot_product_attention: Callable