


'''simple docstring'''



import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin





class    UpperCAmelCase_  (_UpperCAmelCase					,		_UpperCAmelCase  ):





    """simple docstring"""






    @register_to_config
    def __init__(							self			,     SCREAMING_SNAKE_CASE_ = 1_28			,     SCREAMING_SNAKE_CASE_ = 2_56			,     SCREAMING_SNAKE_CASE_ = 2_0_0_0.0			,     SCREAMING_SNAKE_CASE_ = 7_68			,     SCREAMING_SNAKE_CASE_ = 12			,     SCREAMING_SNAKE_CASE_ = 12			,     SCREAMING_SNAKE_CASE_ = 64			,     SCREAMING_SNAKE_CASE_ = 20_48			,     SCREAMING_SNAKE_CASE_ = 0.1			,     )				->     Any:
         super().__init__()

         __lowerCamelCase   :	List[str]								=		nn.Sequential(
             nn.Linear(SCREAMING_SNAKE_CASE_			,     d_model * 4			,     bias=SCREAMING_SNAKE_CASE_	)			,     nn.SiLU()			,     nn.Linear(d_model * 4			,     d_model * 4			,     bias=SCREAMING_SNAKE_CASE_	)			,     nn.SiLU()			,     )

         __lowerCamelCase   :	List[str]								=		nn.Embedding(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		False

         __lowerCamelCase   :	Union[str, Any]								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     bias=SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	str								=		nn.Dropout(p=SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	int								=		nn.ModuleList()
         for lyr_num in range(SCREAMING_SNAKE_CASE_	):
              # FiLM conditional T5 decoder
              __lowerCamelCase   :	List[str]								=		DecoderLayer(d_model=SCREAMING_SNAKE_CASE_			,     d_kv=SCREAMING_SNAKE_CASE_			,     num_heads=SCREAMING_SNAKE_CASE_			,     d_ff=SCREAMING_SNAKE_CASE_			,     dropout_rate=SCREAMING_SNAKE_CASE_	)
              self.decoders.append(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	List[Any]								=		TaLayerNorm(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	str								=		nn.Dropout(p=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	int								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     bias=SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Optional[Any]:
         __lowerCamelCase   :	List[Any]								=		torch.mul(query_input.unsqueeze(-1	)			,     key_input.unsqueeze(-2	)	)
         return mask.unsqueeze(-3	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Optional[Any]:
         __lowerCamelCase				,       __lowerCamelCase				,       __lowerCamelCase   :	Dict								=		decoder_input_tokens.shape
         assert decoder_noise_time.shape == (batch,)

         # decoder_noise_time is in [0, 1), so rescale to expected timing range.
         __lowerCamelCase   :	Dict								=		get_timestep_embedding(
             decoder_noise_time * self.config.max_decoder_noise_time			,     embedding_dim=self.config.d_model			,     max_period=self.config.max_decoder_noise_time			,     ).to(dtype=self.dtype	)

         __lowerCamelCase   :	Union[str, Any]								=		self.conditioning_emb(SCREAMING_SNAKE_CASE_	).unsqueeze(1	)

         assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

         __lowerCamelCase   :	Optional[int]								=		decoder_input_tokens.shape[1]

         # If we want to use relative positions for audio context, we can just offset
         # this sequence by the length of encodings_and_masks.
         __lowerCamelCase   :	Optional[Any]								=		torch.broadcast_to(
             torch.arange(SCREAMING_SNAKE_CASE_			,     device=decoder_input_tokens.device	)			,     (batch, seq_length)			,     )

         __lowerCamelCase   :	Optional[Any]								=		self.position_encoding(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Optional[Any]								=		self.continuous_inputs_projection(SCREAMING_SNAKE_CASE_	)
         inputs += position_encodings
         __lowerCamelCase   :	List[Any]								=		self.dropout(SCREAMING_SNAKE_CASE_	)

         # decoder: No padding present.
         __lowerCamelCase   :	List[Any]								=		torch.ones(
             decoder_input_tokens.shape[:2]			,     device=decoder_input_tokens.device			,     dtype=inputs.dtype	)

         # Translate encoding masks to encoder-decoder masks.
         __lowerCamelCase   :	Dict								=		[(x, self.encoder_decoder_mask(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)) for x, y in encodings_and_masks]

         # cross attend style: concat encodings
         __lowerCamelCase   :	Optional[int]								=		torch.cat([x[0] for x in encodings_and_encdec_masks]			,     dim=1	)
         __lowerCamelCase   :	List[Any]								=		torch.cat([x[1] for x in encodings_and_encdec_masks]			,     dim=-1	)

         for lyr in self.decoders:
              __lowerCamelCase   :	Optional[Any]								=		lyr(
                  SCREAMING_SNAKE_CASE_			,     conditioning_emb=SCREAMING_SNAKE_CASE_			,     encoder_hidden_states=SCREAMING_SNAKE_CASE_			,     encoder_attention_mask=SCREAMING_SNAKE_CASE_			,     )[0]

         __lowerCamelCase   :	Optional[Any]								=		self.decoder_norm(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		self.post_dropout(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	str								=		self.spec_out(SCREAMING_SNAKE_CASE_	)
         return spec_out







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=1E-6	)				->     Dict:
         super().__init__()
         __lowerCamelCase   :	Optional[int]								=		nn.ModuleList()

         # cond self attention: layer 0
         self.layer.append(
             TaLayerSelfAttentionCond(d_model=SCREAMING_SNAKE_CASE_			,     d_kv=SCREAMING_SNAKE_CASE_			,     num_heads=SCREAMING_SNAKE_CASE_			,     dropout_rate=SCREAMING_SNAKE_CASE_	)	)

         # cross attention: layer 1
         self.layer.append(
             TaLayerCrossAttention(
                 d_model=SCREAMING_SNAKE_CASE_			,     d_kv=SCREAMING_SNAKE_CASE_			,     num_heads=SCREAMING_SNAKE_CASE_			,     dropout_rate=SCREAMING_SNAKE_CASE_			,     layer_norm_epsilon=SCREAMING_SNAKE_CASE_			,     )	)

         # Film Cond MLP + dropout: last layer
         self.layer.append(
             TaLayerFFCond(d_model=SCREAMING_SNAKE_CASE_			,     d_ff=SCREAMING_SNAKE_CASE_			,     dropout_rate=SCREAMING_SNAKE_CASE_			,     layer_norm_epsilon=SCREAMING_SNAKE_CASE_	)	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=None			,     SCREAMING_SNAKE_CASE_=None			,     SCREAMING_SNAKE_CASE_=None			,     SCREAMING_SNAKE_CASE_=None			,     SCREAMING_SNAKE_CASE_=None			,     )				->     Dict:
         __lowerCamelCase   :	Optional[Any]								=		self.layer[0](
             SCREAMING_SNAKE_CASE_			,     conditioning_emb=SCREAMING_SNAKE_CASE_			,     attention_mask=SCREAMING_SNAKE_CASE_			,     )

         if encoder_hidden_states is not None:
              __lowerCamelCase   :	Dict								=		torch.where(encoder_attention_mask > 0			,     0			,     -1E10	).to(
                  encoder_hidden_states.dtype	)

              __lowerCamelCase   :	str								=		self.layer[1](
                  SCREAMING_SNAKE_CASE_			,     key_value_states=SCREAMING_SNAKE_CASE_			,     attention_mask=SCREAMING_SNAKE_CASE_			,     )

         # Apply Film Conditional Feed Forward layer
         __lowerCamelCase   :	Union[str, Any]								=		self.layer[-1](SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         return (hidden_states,)







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Dict:
         super().__init__()
         __lowerCamelCase   :	Optional[int]								=		TaLayerNorm(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[int]								=		TaFiLMLayer(in_features=d_model * 4			,     out_features=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		Attention(query_dim=SCREAMING_SNAKE_CASE_			,     heads=SCREAMING_SNAKE_CASE_			,     dim_head=SCREAMING_SNAKE_CASE_			,     out_bias=SCREAMING_SNAKE_CASE_			,     scale_qk=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		nn.Dropout(SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=None			,     SCREAMING_SNAKE_CASE_=None			,     )				->     Tuple:
         # pre_self_attention_layer_norm
         __lowerCamelCase   :	Any								=		self.layer_norm(SCREAMING_SNAKE_CASE_	)

         if conditioning_emb is not None:
              __lowerCamelCase   :	Union[str, Any]								=		self.FiLMLayer(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         # Self-attention block
         __lowerCamelCase   :	int								=		self.attention(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Union[str, Any]								=		hidden_states + self.dropout(SCREAMING_SNAKE_CASE_	)

         return hidden_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Union[str, Any]:
         super().__init__()
         __lowerCamelCase   :	List[str]								=		Attention(query_dim=SCREAMING_SNAKE_CASE_			,     heads=SCREAMING_SNAKE_CASE_			,     dim_head=SCREAMING_SNAKE_CASE_			,     out_bias=SCREAMING_SNAKE_CASE_			,     scale_qk=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		TaLayerNorm(SCREAMING_SNAKE_CASE_			,     eps=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Union[str, Any]								=		nn.Dropout(SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=None			,     SCREAMING_SNAKE_CASE_=None			,     )				->     Optional[Any]:
         __lowerCamelCase   :	str								=		self.layer_norm(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	int								=		self.attention(
             SCREAMING_SNAKE_CASE_			,     encoder_hidden_states=SCREAMING_SNAKE_CASE_			,     attention_mask=attention_mask.squeeze(1	)			,     )
         __lowerCamelCase   :	Optional[Any]								=		hidden_states + self.dropout(SCREAMING_SNAKE_CASE_	)
         return layer_output







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Optional[int]:
         super().__init__()
         __lowerCamelCase   :	Tuple								=		TaDenseGatedActDense(d_model=SCREAMING_SNAKE_CASE_			,     d_ff=SCREAMING_SNAKE_CASE_			,     dropout_rate=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Dict								=		TaFiLMLayer(in_features=d_model * 4			,     out_features=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Dict								=		TaLayerNorm(SCREAMING_SNAKE_CASE_			,     eps=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[Any]								=		nn.Dropout(SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=None	)				->     Tuple:
         __lowerCamelCase   :	str								=		self.layer_norm(SCREAMING_SNAKE_CASE_	)
         if conditioning_emb is not None:
              __lowerCamelCase   :	Tuple								=		self.film(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Any								=		self.DenseReluDense(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		hidden_states + self.dropout(SCREAMING_SNAKE_CASE_	)
         return hidden_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     str:
         super().__init__()
         __lowerCamelCase   :	Optional[Any]								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     bias=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[str]								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     bias=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     bias=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[int]								=		nn.Dropout(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Union[str, Any]								=		NewGELUActivation()







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     List[str]:
         __lowerCamelCase   :	List[str]								=		self.act(self.wi_a(SCREAMING_SNAKE_CASE_	)	)
         __lowerCamelCase   :	int								=		self.wi_a(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[Any]								=		hidden_gelu * hidden_linear
         __lowerCamelCase   :	Optional[Any]								=		self.dropout(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Tuple								=		self.wo(SCREAMING_SNAKE_CASE_	)
         return hidden_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=1E-6	)				->     Tuple:
         super().__init__()
         __lowerCamelCase   :	List[Any]								=		nn.Parameter(torch.ones(SCREAMING_SNAKE_CASE_	)	)
         __lowerCamelCase   :	Any								=		eps







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     Union[str, Any]:
         # T5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean
         # Square Layer Normalization https://arxiv.org/abs/1910.07467 thus variance is calculated
         # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for
         # half-precision inputs is done in fp32

         __lowerCamelCase   :	List[str]								=		hidden_states.to(torch.floataa	).pow(2	).mean(-1			,     keepdim=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		hidden_states * torch.rsqrt(variance + self.variance_epsilon	)

         # convert into half-precision if necessary
         if self.weight.dtype in [torch.floataa, torch.bfloataa]:
              __lowerCamelCase   :	Dict								=		hidden_states.to(self.weight.dtype	)

         return self.weight * hidden_states







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     torch.Tensor:
         return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi	) * (input + 0.0_4_4_7_1_5 * torch.pow(SCREAMING_SNAKE_CASE_			,     3.0	))	))







class    UpperCAmelCase_  (nn.Module  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     int:
         super().__init__()
         __lowerCamelCase   :	Union[str, Any]								=		nn.Linear(SCREAMING_SNAKE_CASE_			,     out_features * 2			,     bias=SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase   :	str								=		self.scale_bias(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase				,       __lowerCamelCase   :	Any								=		torch.chunk(SCREAMING_SNAKE_CASE_			,     2			,     -1	)
         __lowerCamelCase   :	Dict								=		x * (1 + scale) + shift
         return x










'''simple docstring'''



from . import (
    albert,
    align,
    altclip,
    audio_spectrogram_transformer,
    auto,
    autoformer,
    bark,
    bart,
    barthez,
    bartpho,
    beit,
    bert,
    bert_generation,
    bert_japanese,
    bertweet,
    big_bird,
    bigbird_pegasus,
    biogpt,
    bit,
    blenderbot,
    blenderbot_small,
    blip,
    blip_a,
    bloom,
    bridgetower,
    byta,
    camembert,
    canine,
    chinese_clip,
    clap,
    clip,
    clipseg,
    codegen,
    conditional_detr,
    convbert,
    convnext,
    convnextva,
    cpm,
    cpmant,
    ctrl,
    cvt,
    dataavec,
    deberta,
    deberta_va,
    decision_transformer,
    deformable_detr,
    deit,
    deprecated,
    deta,
    detr,
    dialogpt,
    dinat,
    distilbert,
    dit,
    donut,
    dpr,
    dpt,
    efficientformer,
    efficientnet,
    electra,
    encodec,
    encoder_decoder,
    ernie,
    ernie_m,
    esm,
    falcon,
    flaubert,
    flava,
    fnet,
    focalnet,
    fsmt,
    funnel,
    git,
    glpn,
    gpta,
    gpt_bigcode,
    gpt_neo,
    gpt_neox,
    gpt_neox_japanese,
    gpt_swa,
    gptj,
    gptsan_japanese,
    graphormer,
    groupvit,
    herbert,
    hubert,
    ibert,
    imagegpt,
    informer,
    instructblip,
    jukebox,
    layoutlm,
    layoutlmva,
    layoutlmva,
    layoutxlm,
    led,
    levit,
    lilt,
    llama,
    longformer,
    longta,
    luke,
    lxmert,
    mam_aaa,
    marian,
    markuplm,
    maskaformer,
    maskformer,
    mbart,
    mbartaa,
    mega,
    megatron_bert,
    megatron_gpta,
    mgp_str,
    mluke,
    mobilebert,
    mobilenet_va,
    mobilenet_va,
    mobilevit,
    mobilevitva,
    mpnet,
    mra,
    mta,
    musicgen,
    mvp,
    nat,
    nezha,
    nllb,
    nllb_moe,
    nystromformer,
    oneformer,
    open_llama,
    openai,
    opt,
    owlvit,
    pegasus,
    pegasus_x,
    perceiver,
    phobert,
    pixastruct,
    plbart,
    poolformer,
    prophetnet,
    qdqbert,
    rag,
    realm,
    reformer,
    regnet,
    rembert,
    resnet,
    roberta,
    roberta_prelayernorm,
    roc_bert,
    roformer,
    rwkv,
    sam,
    segformer,
    sew,
    sew_d,
    speech_encoder_decoder,
    speech_to_text,
    speech_to_text_a,
    speechta,
    splinter,
    squeezebert,
    swiftformer,
    swin,
    swinasr,
    swinva,
    switch_transformers,
    ta,
    table_transformer,
    tapas,
    time_series_transformer,
    timesformer,
    timm_backbone,
    transfo_xl,
    trocr,
    tvlt,
    umta,
    unispeech,
    unispeech_sat,
    upernet,
    videomae,
    vilt,
    vision_encoder_decoder,
    vision_text_dual_encoder,
    visual_bert,
    vit,
    vit_hybrid,
    vit_mae,
    vit_msn,
    vivit,
    wavaveca,
    wavaveca_conformer,
    wavaveca_phoneme,
    wavaveca_with_lm,
    wavlm,
    whisper,
    x_clip,
    xglm,
    xlm,
    xlm_prophetnet,
    xlm_roberta,
    xlm_roberta_xl,
    xlnet,
    xmod,
    yolos,
    yoso,
)






