from __future__ import annotations

from collections import namedtuple
from typing import Any, Optional

from ... import tensor as tensor
from ...fluid import layers as layers
from ...framework import ParamAttr as ParamAttr
from .. import Layer as Layer
from .. import LayerList as LayerList
from .common import Dropout as Dropout
from .common import Linear as Linear
from .norm import LayerNorm as LayerNorm

class MultiHeadAttention(Layer):
    Cache = namedtuple("Cache", ["k", "v"])

    StaticCache = namedtuple("StaticCache", ["k", "v"])
    embed_dim: Any = ...
    kdim: Any = ...
    vdim: Any = ...
    num_heads: Any = ...
    dropout: Any = ...
    need_weights: Any = ...
    head_dim: Any = ...
    q_proj: Any = ...
    k_proj: Any = ...
    v_proj: Any = ...
    out_proj: Any = ...
    def __init__(
        self,
        embed_dim: Any,
        num_heads: Any,
        dropout: float = ...,
        kdim: Optional[Any] = ...,
        vdim: Optional[Any] = ...,
        need_weights: bool = ...,
        weight_attr: Optional[Any] = ...,
        bias_attr: Optional[Any] = ...,
    ) -> None: ...
    def compute_kv(self, key: Any, value: Any): ...
    def gen_cache(self, key: Any, value: Optional[Any] = ..., type: Any = ...): ...
    def forward(
        self,
        query: Any,
        key: Optional[Any] = ...,
        value: Optional[Any] = ...,
        attn_mask: Optional[Any] = ...,
        cache: Optional[Any] = ...,
    ): ...

class TransformerEncoderLayer(Layer):
    normalize_before: Any = ...
    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: Any,
        dropout: float = ...,
        activation: str = ...,
        attn_dropout: Optional[Any] = ...,
        act_dropout: Optional[Any] = ...,
        normalize_before: bool = ...,
        weight_attr: Optional[Any] = ...,
        bias_attr: Optional[Any] = ...,
    ) -> None: ...
    def forward(self, src: Any, src_mask: Optional[Any] = ..., cache: Optional[Any] = ...): ...
    def gen_cache(self, src: Any): ...

class TransformerEncoder(Layer):
    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, src_mask: Optional[Any] = ..., cache: Optional[Any] = ...): ...
    def gen_cache(self, src: Any): ...

class TransformerDecoderLayer(Layer):
    normalize_before: Any = ...
    self_attn: Any = ...
    cross_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: Any,
        dropout: float = ...,
        activation: str = ...,
        attn_dropout: Optional[Any] = ...,
        act_dropout: Optional[Any] = ...,
        normalize_before: bool = ...,
        weight_attr: Optional[Any] = ...,
        bias_attr: Optional[Any] = ...,
    ) -> None: ...
    def forward(
        self,
        tgt: Any,
        memory: Any,
        tgt_mask: Optional[Any] = ...,
        memory_mask: Optional[Any] = ...,
        cache: Optional[Any] = ...,
    ): ...
    def gen_cache(self, memory: Any): ...

class TransformerDecoder(Layer):
    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] = ...,
        cache: Optional[Any] = ...,
    ): ...
    def gen_cache(self, memory: Any, do_zip: bool = ...): ...

class Transformer(Layer):
    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 = ...,
        attn_dropout: Optional[Any] = ...,
        act_dropout: Optional[Any] = ...,
        normalize_before: bool = ...,
        weight_attr: Optional[Any] = ...,
        bias_attr: Optional[Any] = ...,
        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] = ...,
    ): ...
    def generate_square_subsequent_mask(self, length: Any): ...
