torch/nn/modules/transformer.pyi.in

from ..init import xavier_uniform_
from .activation import MultiheadAttention
from .container import ModuleList
from .dropout import Dropout
from .linear import Linear
from .module import Module
from .normalization import LayerNorm
from typing import Any, Optional

class Transformer(Module):
    encoder: Any = ...
    decoder: Any = ...
    d_model: Any = ...
    nhead: Any = ...
    def __init__(self, d_model: int = ..., nhead: int = ..., num_encoder_layers: int = ..., num_decoder_layers: int = ..., dim_feedforward: int = ..., dropout: float = ..., activation: str = ..., custom_encoder: Optional[Any] = ..., custom_decoder: Optional[Any] = ...) -> None: ...
    def forward(self, src: Any, tgt: Any, src_mask: Optional[Any] = ..., tgt_mask: Optional[Any] = ..., memory_mask: Optional[Any] = ..., src_key_padding_mask: Optional[Any] = ..., tgt_key_padding_mask: Optional[Any] = ..., memory_key_padding_mask: Optional[Any] = ...): ...
    def generate_square_subsequent_mask(self, sz: Any): ...

class TransformerEncoder(Module):
    layers: Any = ...
    num_layers: Any = ...
    norm: Any = ...
    def __init__(self, encoder_layer: Any, num_layers: Any, norm: Optional[Any] = ...) -> None: ...
    def forward(self, src: Any, mask: Optional[Any] = ..., src_key_padding_mask: Optional[Any] = ...): ...

class TransformerDecoder(Module):
    layers: Any = ...
    num_layers: Any = ...
    norm: Any = ...
    def __init__(self, decoder_layer: Any, num_layers: Any, norm: Optional[Any] = ...) -> None: ...
    def forward(self, tgt: Any, memory: Any, tgt_mask: Optional[Any] = ..., memory_mask: Optional[Any] = ..., tgt_key_padding_mask: Optional[Any] = ..., memory_key_padding_mask: Optional[Any] = ...): ...

class TransformerEncoderLayer(Module):
    self_attn: Any = ...
    linear1: Any = ...
    dropout: Any = ...
    linear2: Any = ...
    norm1: Any = ...
    norm2: Any = ...
    dropout1: Any = ...
    dropout2: Any = ...
    activation: Any = ...
    def __init__(self, d_model: Any, nhead: Any, dim_feedforward: int = ..., dropout: float = ..., activation: str = ...) -> None: ...
    def forward(self, src: Any, src_mask: Optional[Any] = ..., src_key_padding_mask: Optional[Any] = ...): ...

class TransformerDecoderLayer(Module):
    self_attn: Any = ...
    multihead_attn: Any = ...
    linear1: Any = ...
    dropout: Any = ...
    linear2: Any = ...
    norm1: Any = ...
    norm2: Any = ...
    norm3: Any = ...
    dropout1: Any = ...
    dropout2: Any = ...
    dropout3: Any = ...
    activation: Any = ...
    def __init__(self, d_model: Any, nhead: Any, dim_feedforward: int = ..., dropout: float = ..., activation: str = ...) -> None: ...
    def forward(self, tgt: Any, memory: Any, tgt_mask: Optional[Any] = ..., memory_mask: Optional[Any] = ..., tgt_key_padding_mask: Optional[Any] = ..., memory_key_padding_mask: Optional[Any] = ...): ...