from __future__ import annotations

from typing import Any, Optional

from paddle.fluid import dygraph_utils as dygraph_utils

def fused_feedforward(
    x: Any,
    linear1_weight: Any,
    linear2_weight: Any,
    linear1_bias: Optional[Any] = ...,
    linear2_bias: Optional[Any] = ...,
    ln1_scale: Optional[Any] = ...,
    ln1_bias: Optional[Any] = ...,
    ln2_scale: Optional[Any] = ...,
    ln2_bias: Optional[Any] = ...,
    dropout1_rate: float = ...,
    dropout2_rate: float = ...,
    activation: str = ...,
    ln1_epsilon: float = ...,
    ln2_epsilon: float = ...,
    pre_layer_norm: bool = ...,
    training: bool = ...,
    mode: str = ...,
    ring_id: int = ...,
    add_residual: bool = ...,
    name: Optional[Any] = ...,
): ...
def fused_multi_head_attention(
    x: Any,
    qkv_weight: Any,
    linear_weight: Any,
    pre_layer_norm: bool = ...,
    pre_ln_scale: Optional[Any] = ...,
    pre_ln_bias: Optional[Any] = ...,
    ln_scale: Optional[Any] = ...,
    ln_bias: Optional[Any] = ...,
    pre_ln_epsilon: float = ...,
    qkv_bias: Optional[Any] = ...,
    linear_bias: Optional[Any] = ...,
    cache_kv: Optional[Any] = ...,
    attn_mask: Optional[Any] = ...,
    dropout_rate: float = ...,
    attn_dropout_rate: float = ...,
    ln_epsilon: float = ...,
    training: bool = ...,
    mode: str = ...,
    ring_id: int = ...,
    add_residual: bool = ...,
    name: Optional[Any] = ...,
): ...
def fused_multi_transformer(
    x: Any,
    ln_scales: Any,
    ln_biases: Any,
    qkv_weights: Any,
    qkv_biases: Any,
    linear_weights: Any,
    linear_biases: Any,
    ffn_ln_scales: Any,
    ffn_ln_biases: Any,
    ffn1_weights: Any,
    ffn1_biases: Any,
    ffn2_weights: Any,
    ffn2_biases: Any,
    pre_layer_norm: bool = ...,
    epsilon: float = ...,
    cache_kvs: Optional[Any] = ...,
    time_step: Optional[Any] = ...,
    attn_mask: Optional[Any] = ...,
    dropout_rate: float = ...,
    activation: str = ...,
    training: bool = ...,
    mode: str = ...,
    ring_id: int = ...,
    name: Optional[Any] = ...,
): ...
