




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__      (SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_	):

    '''simple docstring'''



    @register_to_config
    def __init__(       self	,    UpperCamelCase__ = 128	,    UpperCamelCase__ = 256	,    UpperCamelCase__ = 2000.0	,    UpperCamelCase__ = 768	,    UpperCamelCase__ = 12	,    UpperCamelCase__ = 12	,    UpperCamelCase__ = 64	,    UpperCamelCase__ = 2048	,    UpperCamelCase__ = 0.1	,    )	-> List[Any]:
           super().__init__()

           lowerCamelCase  : int     					=   nn.Sequential(
               nn.Linear(UpperCamelCase__	,    d_model * 4	,    bias=UpperCamelCase__	)	,    nn.SiLU()	,    nn.Linear(d_model * 4	,    d_model * 4	,    bias=UpperCamelCase__	)	,    nn.SiLU()	,    )

           lowerCamelCase  : Any     					=   nn.Embedding(UpperCamelCase__	,    UpperCamelCase__	)
           lowerCamelCase  : str     					=   False

           lowerCamelCase  : Any     					=   nn.Linear(UpperCamelCase__	,    UpperCamelCase__	,    bias=UpperCamelCase__	)

           lowerCamelCase  : Any     					=   nn.Dropout(p=UpperCamelCase__	)

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

           lowerCamelCase  : Union[str, Any]     					=   TaLayerNorm(UpperCamelCase__	)

           lowerCamelCase  : Dict     					=   nn.Dropout(p=UpperCamelCase__	)
           lowerCamelCase  : int     					=   nn.Linear(UpperCamelCase__	,    UpperCamelCase__	,    bias=UpperCamelCase__	)







    def        _lowercase	(       self	,    UpperCamelCase__	,    UpperCamelCase__	)	-> int:
           lowerCamelCase  : str     					=   torch.mul(query_input.unsqueeze(-1	)	,    key_input.unsqueeze(-2	)	)
           return mask.unsqueeze(-3	)







    def        _lowercase	(       self	,    UpperCamelCase__	,    UpperCamelCase__	,    UpperCamelCase__	)	-> Union[str, Any]:
           lowerCamelCase  : List[Any]     					=   decoder_input_tokens.shape
           assert decoder_noise_time.shape == (batch,)

           # decoder_noise_time is in [0, 1), so rescale to expected timing range.
           lowerCamelCase  : Any     					=   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  : str     					=   self.conditioning_emb(UpperCamelCase__	).unsqueeze(1	)

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

           lowerCamelCase  : Optional[Any]     					=   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  : List[Any]     					=   torch.broadcast_to(
               torch.arange(UpperCamelCase__	,    device=decoder_input_tokens.device	)	,    (batch, seq_length)	,    )

           lowerCamelCase  : int     					=   self.position_encoding(UpperCamelCase__	)

           lowerCamelCase  : str     					=   self.continuous_inputs_projection(UpperCamelCase__	)
           inputs += position_encodings
           lowerCamelCase  : List[Any]     					=   self.dropout(UpperCamelCase__	)

           # decoder: No padding present.
           lowerCamelCase  : Union[str, 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(UpperCamelCase__	,    UpperCamelCase__	)) 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  : Union[str, Any]     					=   torch.cat([x[1] for x in encodings_and_encdec_masks]	,    dim=-1	)

           for lyr in self.decoders:
                  lowerCamelCase  : int     					=   lyr(
                      UpperCamelCase__	,    conditioning_emb=UpperCamelCase__	,    encoder_hidden_states=UpperCamelCase__	,    encoder_attention_mask=UpperCamelCase__	,    )[0]

           lowerCamelCase  : Dict     					=   self.decoder_norm(UpperCamelCase__	)
           lowerCamelCase  : Optional[Any]     					=   self.post_dropout(UpperCamelCase__	)

           lowerCamelCase  : Dict     					=   self.spec_out(UpperCamelCase__	)
           return spec_out




class 						UpperCamelCase__      (nn.Module	):

    '''simple docstring'''



    def __init__(       self	,    UpperCamelCase__	,    UpperCamelCase__	,    UpperCamelCase__	,    UpperCamelCase__	,    UpperCamelCase__	,    UpperCamelCase__=1e-6	)	-> Optional[int]:
           super().__init__()
           lowerCamelCase  : str     					=   nn.ModuleList()

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

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

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







    def        _lowercase	(       self	,    UpperCamelCase__	,    UpperCamelCase__=None	,    UpperCamelCase__=None	,    UpperCamelCase__=None	,    UpperCamelCase__=None	,    UpperCamelCase__=None	,    )	-> List[str]:
           lowerCamelCase  : Optional[int]     					=   self.layer[0](
               UpperCamelCase__	,    conditioning_emb=UpperCamelCase__	,    attention_mask=UpperCamelCase__	,    )

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

                  lowerCamelCase  : Any     					=   self.layer[1](
                      UpperCamelCase__	,    key_value_states=UpperCamelCase__	,    attention_mask=UpperCamelCase__	,    )

           # Apply Film Conditional Feed Forward layer
           lowerCamelCase  : Dict     					=   self.layer[-1](UpperCamelCase__	,    UpperCamelCase__	)

           return (hidden_states,)




class 						UpperCamelCase__      (nn.Module	):

    '''simple docstring'''



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







    def        _lowercase	(       self	,    UpperCamelCase__	,    UpperCamelCase__=None	,    UpperCamelCase__=None	,    )	-> Dict:
           lowerCamelCase  : Tuple     					=   self.layer_norm(UpperCamelCase__	)

           if conditioning_emb is not None:
                  lowerCamelCase  : Dict     					=   self.FiLMLayer(UpperCamelCase__	,    UpperCamelCase__	)

           # Self-attention block
           lowerCamelCase  : str     					=   self.attention(UpperCamelCase__	)

           lowerCamelCase  : Dict     					=   hidden_states + self.dropout(UpperCamelCase__	)

           return hidden_states




class 						UpperCamelCase__      (nn.Module	):

    '''simple docstring'''



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







    def        _lowercase	(       self	,    UpperCamelCase__	,    UpperCamelCase__=None	,    UpperCamelCase__=None	,    )	-> Optional[int]:
           lowerCamelCase  : str     					=   self.layer_norm(UpperCamelCase__	)
           lowerCamelCase  : Dict     					=   self.attention(
               UpperCamelCase__	,    encoder_hidden_states=UpperCamelCase__	,    attention_mask=attention_mask.squeeze(1	)	,    )
           lowerCamelCase  : List[str]     					=   hidden_states + self.dropout(UpperCamelCase__	)
           return layer_output




class 						UpperCamelCase__      (nn.Module	):

    '''simple docstring'''



    def __init__(       self	,    UpperCamelCase__	,    UpperCamelCase__	,    UpperCamelCase__	,    UpperCamelCase__	)	-> Dict:
           super().__init__()
           lowerCamelCase  : Dict     					=   TaDenseGatedActDense(d_model=UpperCamelCase__	,    d_ff=UpperCamelCase__	,    dropout_rate=UpperCamelCase__	)
           lowerCamelCase  : List[str]     					=   TaFiLMLayer(in_features=d_model * 4	,    out_features=UpperCamelCase__	)
           lowerCamelCase  : str     					=   TaLayerNorm(UpperCamelCase__	,    eps=UpperCamelCase__	)
           lowerCamelCase  : Union[str, Any]     					=   nn.Dropout(UpperCamelCase__	)







    def        _lowercase	(       self	,    UpperCamelCase__	,    UpperCamelCase__=None	)	-> int:
           lowerCamelCase  : Tuple     					=   self.layer_norm(UpperCamelCase__	)
           if conditioning_emb is not None:
                  lowerCamelCase  : Union[str, Any]     					=   self.film(UpperCamelCase__	,    UpperCamelCase__	)

           lowerCamelCase  : Any     					=   self.DenseReluDense(UpperCamelCase__	)
           lowerCamelCase  : List[Any]     					=   hidden_states + self.dropout(UpperCamelCase__	)
           return hidden_states




class 						UpperCamelCase__      (nn.Module	):

    '''simple docstring'''



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







    def        _lowercase	(       self	,    UpperCamelCase__	)	-> int:
           lowerCamelCase  : Tuple     					=   self.act(self.wi_a(UpperCamelCase__	)	)
           lowerCamelCase  : Any     					=   self.wi_a(UpperCamelCase__	)
           lowerCamelCase  : Any     					=   hidden_gelu * hidden_linear
           lowerCamelCase  : Union[str, Any]     					=   self.dropout(UpperCamelCase__	)

           lowerCamelCase  : int     					=   self.wo(UpperCamelCase__	)
           return hidden_states




class 						UpperCamelCase__      (nn.Module	):

    '''simple docstring'''



    def __init__(       self	,    UpperCamelCase__	,    UpperCamelCase__=1e-6	)	-> List[str]:
           super().__init__()
           lowerCamelCase  : Optional[Any]     					=   nn.Parameter(torch.ones(UpperCamelCase__	)	)
           lowerCamelCase  : Optional[int]     					=   eps







    def        _lowercase	(       self	,    UpperCamelCase__	)	-> List[str]:
           lowerCamelCase  : List[str]     					=   hidden_states.to(torch.floataa	).pow(2	).mean(-1	,    keepdim=UpperCamelCase__	)
           lowerCamelCase  : List[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  : List[Any]     					=   hidden_states.to(self.weight.dtype	)

           return self.weight * hidden_states




class 						UpperCamelCase__      (nn.Module	):

    '''simple docstring'''



    def        _lowercase	(       self	,    UpperCamelCase__	)	-> int:
           return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi	) * (input + 0.044715 * torch.pow(UpperCamelCase__	,    3.0	))	))




class 						UpperCamelCase__      (nn.Module	):

    '''simple docstring'''



    def __init__(       self	,    UpperCamelCase__	,    UpperCamelCase__	)	-> Dict:
           super().__init__()
           lowerCamelCase  : int     					=   nn.Linear(UpperCamelCase__	,    out_features * 2	,    bias=UpperCamelCase__	)







    def        _lowercase	(       self	,    UpperCamelCase__	,    UpperCamelCase__	)	-> str:
           lowerCamelCase  : Tuple     					=   self.scale_bias(UpperCamelCase__	)
           lowerCamelCase  : Optional[int]     					=   torch.chunk(UpperCamelCase__	,    2	,    -1	)
           lowerCamelCase  : Optional[int]     					=   x * (1 + scale) + shift
           return x











import unittest

import numpy as np

from transformers import DistilBertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
			import jax.numpy as jnp

			from transformers.models.distilbert.modeling_flax_distilbert import (
			    FlaxDistilBertForMaskedLM,
			    FlaxDistilBertForMultipleChoice,
			    FlaxDistilBertForQuestionAnswering,
			    FlaxDistilBertForSequenceClassification,
			    FlaxDistilBertForTokenClassification,
			    FlaxDistilBertModel,
			)


class 						lowercase    (							unittest.TestCase):
						'''simple docstring'''


						def __init__(  self    :				int       ,	snake_case    :				Optional[Any]       ,	snake_case    :				Optional[Any]=13       ,	snake_case    :				str=7       ,	snake_case    :				Optional[Any]=True       ,	snake_case    :				int=True       ,	snake_case    :				str=True       ,	snake_case    :				List[str]=True       ,	snake_case    :				Optional[Any]=99       ,	snake_case    :				Optional[int]=32       ,	snake_case    :				Optional[int]=5       ,	snake_case    :				Optional[Any]=4       ,	snake_case    :				Optional[Any]=37       ,	snake_case    :				str="gelu"       ,	snake_case    :				List[str]=0.1       ,	snake_case    :				List[str]=0.1       ,	snake_case    :				int=512       ,	snake_case    :				Union[str, Any]=16       ,	snake_case    :				Optional[Any]=2       ,	snake_case    :				int=0.02       ,	snake_case    :				str=4       ,	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Tuple						=				parent
											SCREAMING_SNAKE_CASE							:     str						=				batch_size
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				seq_length
											SCREAMING_SNAKE_CASE							:     int						=				is_training
											SCREAMING_SNAKE_CASE							:     List[str]						=				use_attention_mask
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				use_token_type_ids
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				use_labels
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				vocab_size
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				hidden_size
											SCREAMING_SNAKE_CASE							:     Any						=				num_hidden_layers
											SCREAMING_SNAKE_CASE							:     Tuple						=				num_attention_heads
											SCREAMING_SNAKE_CASE							:     Any						=				intermediate_size
											SCREAMING_SNAKE_CASE							:     Tuple						=				hidden_act
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				hidden_dropout_prob
											SCREAMING_SNAKE_CASE							:     str						=				attention_probs_dropout_prob
											SCREAMING_SNAKE_CASE							:     Tuple						=				max_position_embeddings
											SCREAMING_SNAKE_CASE							:     Any						=				type_vocab_size
											SCREAMING_SNAKE_CASE							:     List[Any]						=				type_sequence_label_size
											SCREAMING_SNAKE_CASE							:     Tuple						=				initializer_range
											SCREAMING_SNAKE_CASE							:     List[str]						=				num_choices







						def       lowerCamelCase_				(  self    :				List[Any]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				ids_tensor([self.batch_size, self.seq_length]       ,	self.vocab_size	)

											SCREAMING_SNAKE_CASE							:     Optional[int]						=				None
											if self.use_attention_mask:
																SCREAMING_SNAKE_CASE							:     Any						=				random_attention_mask([self.batch_size, self.seq_length]	)

											SCREAMING_SNAKE_CASE							:     Tuple						=				DistilBertConfig(
											    vocab_size=self.vocab_size       ,	dim=self.hidden_size       ,	n_layers=self.num_hidden_layers       ,	n_heads=self.num_attention_heads       ,	hidden_dim=self.intermediate_size       ,	hidden_act=self.hidden_act       ,	dropout=self.hidden_dropout_prob       ,	attention_dropout=self.attention_probs_dropout_prob       ,	max_position_embeddings=self.max_position_embeddings       ,	initializer_range=self.initializer_range       ,	tie_weights_=snake_case       ,	)

											return config, input_ids, attention_mask







						def       lowerCamelCase_				(  self    :				int	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Tuple						=				self.prepare_config_and_inputs()
											SCREAMING_SNAKE_CASE					,  SCREAMING_SNAKE_CASE					,  SCREAMING_SNAKE_CASE							:     Optional[int]						=				config_and_inputs
											SCREAMING_SNAKE_CASE							:     int						=				{'input_ids': input_ids, 'attention_mask': attention_mask}
											return config, inputs_dict


@require_flax
class 						lowercase    (							SCREAMING_SNAKE_CASE_      ,    unittest.TestCase):
						'''simple docstring'''


						UpperCAmelCase			:		int		=		(
						    (
						        FlaxDistilBertModel,
						        FlaxDistilBertForMaskedLM,
						        FlaxDistilBertForMultipleChoice,
						        FlaxDistilBertForQuestionAnswering,
						        FlaxDistilBertForSequenceClassification,
						        FlaxDistilBertForTokenClassification,
						        FlaxDistilBertForQuestionAnswering,
						    )
						    if is_flax_available()
						    else ()
						)







						def       lowerCamelCase_				(  self    :				Optional[int]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     str						=				FlaxDistilBertModelTester(self	)







						@slow
						def       lowerCamelCase_				(  self    :				Any	):



											'''simple docstring'''


											for model_class_name in self.all_model_classes:
																SCREAMING_SNAKE_CASE							:     Dict						=				model_class_name.from_pretrained('distilbert-base-uncased'	)
																SCREAMING_SNAKE_CASE							:     Dict						=				model(np.ones((1, 1)	)	)
																self.assertIsNotNone(snake_case	)


@require_flax
class 						lowercase    (							unittest.TestCase):
						'''simple docstring'''


						@slow
						def       lowerCamelCase_				(  self    :				Tuple	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     List[Any]						=				FlaxDistilBertModel.from_pretrained('distilbert-base-uncased'	)
											SCREAMING_SNAKE_CASE							:     List[str]						=				np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]	)
											SCREAMING_SNAKE_CASE							:     Dict						=				np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]	)
											SCREAMING_SNAKE_CASE							:     str						=				model(snake_case       ,	attention_mask=snake_case	)[0]
											SCREAMING_SNAKE_CASE							:     int						=				(1, 11, 768)
											self.assertEqual(output.shape       ,	snake_case	)
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				np.array([[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]	)

											self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]       ,	snake_case       ,	atol=1E-4	)	)