89 lines
4.2 KiB
Python
89 lines
4.2 KiB
Python
# @generated by tools/pyi/gen_pyi.py from torch/_C/_nn.pyi.in
|
|
# mypy: disable-error-code="type-arg"
|
|
|
|
from typing import Literal, Optional, overload, Sequence, Union
|
|
|
|
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_max_pool2d(input: Tensor, output_size: Union[_int, _size]) -> tuple[Tensor, Tensor]: ...
|
|
def adaptive_max_pool3d(input: Tensor, output_size: Union[_int, _size]) -> tuple[Tensor, Tensor]: ...
|
|
def avg_pool2d(input: Tensor, kernel_size: Union[_int, _size], stride: Optional[Union[_int, _size]] = None, padding: Union[_int, _size] = 0, ceil_mode: bool = False, count_include_pad: bool = True, divisor_override: Optional[int] = None) -> Tensor: ...
|
|
def avg_pool3d(input: Tensor, kernel_size: Union[_int, _size], stride: Optional[Union[_int, _size]] = None, padding: Union[_int, _size] = 0, ceil_mode: bool = False, count_include_pad: bool = True, divisor_override: Optional[int] = None) -> Tensor: ...
|
|
def elu_(input: Tensor, alpha: float = ...) -> Tensor: ...
|
|
def fractional_max_pool2d(input: Tensor, kernel_size: Union[_int, _size], output_size: Union[_int, _size], _random_samples: Tensor) -> tuple[Tensor, Tensor]: ...
|
|
def fractional_max_pool3d(input: Tensor, kernel_size: Union[_int, _size], output_size: Union[_int, _size], _random_samples: Tensor) -> tuple[Tensor, Tensor]: ...
|
|
def gelu(input: Tensor, approximate: str = ...) -> Tensor: ...
|
|
def hardsigmoid(input: Tensor, *, out: Optional[Tensor] = None) -> Tensor: ...
|
|
def hardtanh(input: Tensor, min_val: float = ..., max_val: float = ..., *, out: Optional[Tensor] = None) -> Tensor: ...
|
|
def hardtanh_(input: Tensor, min_val: float = ..., max_val: float = ...) -> Tensor: ...
|
|
def leaky_relu(input: Tensor, negative_slope: float = ..., *, out: Optional[Tensor] = None) -> Tensor: ...
|
|
def leaky_relu_(input: Tensor, negative_slope: float = ...) -> Tensor: ...
|
|
def linear(input: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor: ...
|
|
def log_sigmoid(input: Tensor) -> Tensor: ...
|
|
def one_hot(tensor: Tensor, num_classes: int = ...) -> Tensor: ...
|
|
def pad(input: Tensor, pad: Sequence[int], mode: str = ..., value: Optional[float] = None) -> Tensor: ...
|
|
def scaled_dot_product_attention(query: Tensor, key: Tensor, value: Tensor, attn_mask: Optional[Tensor] = None, dropout_p: float = 0.0, is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False) -> 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: Optional[Tensor]) -> 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: Union[list[Tensor], tuple[Tensor, ...]],
|
|
batch_first: bool = False,
|
|
padding_value: float = 0.0,
|
|
padding_side: Union[Literal["left", "right"], str] = "right",
|
|
) -> Tensor: ...
|
|
def flatten_dense_tensors(tensors: list[Tensor]) -> Tensor: ...
|
|
def unflatten_dense_tensors(flat: Tensor, tensors: list[Tensor]) -> list[Tensor]: ...
|