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] = ...): ...  # type: ignore
    def __call__(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] = ...): ...  # type: ignore
    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] = ...): ...  # type: ignore
    def __call__(self, src: Any, mask: Optional[Any] = ..., src_key_padding_mask: Optional[Any] = ...): ...  # type: ignore

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] = ...): ...  # type: ignore
    def __call__(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] = ...): ...  # type: ignore

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] = ...): ...  # type: ignore
    def __call__(self, src: Any, src_mask: Optional[Any] = ..., src_key_padding_mask: Optional[Any] = ...): ...  # type: ignore

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] = ...): ...  # type: ignore
    def __call__(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] = ...): ...  # type: ignore
