
"""simple docstring"""




import unittest

import numpy as np

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
    from transformers.models.bert.modeling_flax_bert import (
        FlaxBertForMaskedLM,
        FlaxBertForMultipleChoice,
        FlaxBertForNextSentencePrediction,
        FlaxBertForPreTraining,
        FlaxBertForQuestionAnswering,
        FlaxBertForSequenceClassification,
        FlaxBertForTokenClassification,
        FlaxBertModel,
    )







class   _lowercase						(       unittest.TestCase				):


     '''simple docstring'''





     def __init__(      self ,			__UpperCamelCase ,			__UpperCamelCase=13 ,			__UpperCamelCase=7 ,			__UpperCamelCase=True ,			__UpperCamelCase=True ,			__UpperCamelCase=True ,			__UpperCamelCase=True ,			__UpperCamelCase=99 ,			__UpperCamelCase=32 ,			__UpperCamelCase=5 ,			__UpperCamelCase=4 ,			__UpperCamelCase=37 ,			__UpperCamelCase="gelu" ,			__UpperCamelCase=0.1 ,			__UpperCamelCase=0.1 ,			__UpperCamelCase=5_12 ,			__UpperCamelCase=16 ,			__UpperCamelCase=2 ,			__UpperCamelCase=0.02 ,			__UpperCamelCase=4 ,			)->					Optional[int]:

          '''simple docstring'''





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



     def  lowerCAmelCase__       (      self  )->					Optional[Any]:

          '''simple docstring'''





          UpperCAmelCase__							: List[Any] 							=						ids_tensor([self.batch_size, self.seq_length] ,			self.vocab_size  )

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

          UpperCAmelCase__							: Union[str, Any] 							=						None
          if self.use_token_type_ids:
               UpperCAmelCase__							: List[Any] 							=						ids_tensor([self.batch_size, self.seq_length] ,			self.type_vocab_size  )

          UpperCAmelCase__							: Optional[Any] 							=						BertConfig(
              vocab_size=self.vocab_size ,			hidden_size=self.hidden_size ,			num_hidden_layers=self.num_hidden_layers ,			num_attention_heads=self.num_attention_heads ,			intermediate_size=self.intermediate_size ,			hidden_act=self.hidden_act ,			hidden_dropout_prob=self.hidden_dropout_prob ,			attention_probs_dropout_prob=self.attention_probs_dropout_prob ,			max_position_embeddings=self.max_position_embeddings ,			type_vocab_size=self.type_vocab_size ,			is_decoder=_A ,			initializer_range=self.initializer_range ,			)

          return config, input_ids, token_type_ids, attention_mask



     def  lowerCAmelCase__       (      self  )->					Optional[Any]:

          '''simple docstring'''





          UpperCAmelCase__							: List[str] 							=						self.prepare_config_and_inputs()
          UpperCAmelCase__							: str 							=						config_and_inputs
          UpperCAmelCase__							: List[str] 							=						{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
          return config, inputs_dict



     def  lowerCAmelCase__       (      self  )->					Dict:

          '''simple docstring'''





          UpperCAmelCase__							: Optional[int] 							=						self.prepare_config_and_inputs()
          UpperCAmelCase__							: Optional[Any] 							=						config_and_inputs

          UpperCAmelCase__							: Dict 							=						True
          UpperCAmelCase__							: Tuple 							=						floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
          UpperCAmelCase__							: Tuple 							=						ids_tensor([self.batch_size, self.seq_length] ,			vocab_size=2  )

          return (
              config,
              input_ids,
              attention_mask,
              encoder_hidden_states,
              encoder_attention_mask,
          )







@require_flax
class   _lowercase						(       snake_case__	,   unittest.TestCase				):


     '''simple docstring'''





     _A								=							True

     _A								=							(
         (
             FlaxBertModel,
             FlaxBertForPreTraining,
             FlaxBertForMaskedLM,
             FlaxBertForMultipleChoice,
             FlaxBertForQuestionAnswering,
             FlaxBertForNextSentencePrediction,
             FlaxBertForSequenceClassification,
             FlaxBertForTokenClassification,
             FlaxBertForQuestionAnswering,
         )
         if is_flax_available()
         else ()
     )



     def  lowerCAmelCase__       (      self  )->					Optional[int]:

          '''simple docstring'''





          UpperCAmelCase__							: Optional[Any] 							=						FlaxBertModelTester(self  )



     @slow
     def  lowerCAmelCase__       (      self  )->					str:

          '''simple docstring'''





          UpperCAmelCase__							: List[str] 							=						FlaxBertModel.from_pretrained("bert-base-cased"  )
          UpperCAmelCase__							: Tuple 							=						model(np.ones((1, 1)  )  )
          self.assertIsNotNone(_A  )





"""simple docstring"""




from typing import Optional, Union

import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models.modeling_utils import ModelMixin







class   _lowercase						(       lowerCAmelCase_	,   lowerCAmelCase_				):


     '''simple docstring'''





     @register_to_config
     def __init__(      self ,			__UpperCamelCase = 7_68 ,			)->					Union[str, Any]:
          super().__init__()

          UpperCAmelCase__							: str 							=						nn.Parameter(torch.zeros(1 ,			__UpperCamelCase  )  )
          UpperCAmelCase__							: Optional[int] 							=						nn.Parameter(torch.ones(1 ,			__UpperCamelCase  )  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase = None ,			__UpperCamelCase = None ,			)->					Any:
          UpperCAmelCase__							: Dict 							=						nn.Parameter(self.mean.to(__UpperCamelCase  ).to(__UpperCamelCase  )  )
          UpperCAmelCase__							: Any 							=						nn.Parameter(self.std.to(__UpperCamelCase  ).to(__UpperCamelCase  )  )
          return self



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase  )->					Union[str, Any]:
          UpperCAmelCase__							: Dict 							=						(embeds - self.mean) * 1.0 / self.std
          return embeds



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase  )->					List[Any]:
          UpperCAmelCase__							: Any 							=						(embeds * self.std) + self.mean
          return embeds



