





from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase_  							=  logging.get_logger(__name__)

UpperCAmelCase_  							=  {
    # See all MEGATRON_BERT models at https://huggingface.co/models?filter=bert
}






class 					lowercase__ (			__lowerCamelCase       ):

   '''simple docstring'''


   a		:							List[Any]     	=	"megatron-bert"




   def __init__(					self,	__magic_name__=29056,	__magic_name__=1024,	__magic_name__=24,	__magic_name__=16,	__magic_name__=4096,	__magic_name__="gelu",	__magic_name__=0.1,	__magic_name__=0.1,	__magic_name__=512,	__magic_name__=2,	__magic_name__=0.02,	__magic_name__=1E-12,	__magic_name__=0,	__magic_name__="absolute",	__magic_name__=True,	**__magic_name__,	)				->   Optional[int]:






       """simple docstring"""
       super().__init__(pad_token_id=__magic_name__,	**__magic_name__							)

       UpperCamelCase__							:       List[Any]      =   vocab_size
       UpperCamelCase__							:       Optional[Any]      =   hidden_size
       UpperCamelCase__							:       str      =   num_hidden_layers
       UpperCamelCase__							:       int      =   num_attention_heads
       UpperCamelCase__							:       List[Any]      =   hidden_act
       UpperCamelCase__							:       Tuple      =   intermediate_size
       UpperCamelCase__							:       str      =   hidden_dropout_prob
       UpperCamelCase__							:       Any      =   attention_probs_dropout_prob
       UpperCamelCase__							:       Optional[int]      =   max_position_embeddings
       UpperCamelCase__							:       Dict      =   type_vocab_size
       UpperCamelCase__							:       Optional[Any]      =   initializer_range
       UpperCamelCase__							:       Union[str, Any]      =   layer_norm_eps
       UpperCamelCase__							:       Union[str, Any]      =   position_embedding_type
       UpperCamelCase__							:       Optional[Any]      =   use_cache








import unittest

from transformers import BigBirdConfig, 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

  from transformers.models.big_bird.modeling_flax_big_bird import (
      FlaxBigBirdForCausalLM,
      FlaxBigBirdForMaskedLM,
      FlaxBigBirdForMultipleChoice,
      FlaxBigBirdForPreTraining,
      FlaxBigBirdForQuestionAnswering,
      FlaxBigBirdForSequenceClassification,
      FlaxBigBirdForTokenClassification,
      FlaxBigBirdModel,
  )






class 					lowercase__ (			unittest.TestCase       ):

   '''simple docstring'''


   def __init__(					self,	__magic_name__,	__magic_name__=2,	__magic_name__=56,	__magic_name__=True,	__magic_name__=True,	__magic_name__=True,	__magic_name__=True,	__magic_name__=99,	__magic_name__=32,	__magic_name__=2,	__magic_name__=2,	__magic_name__=7,	__magic_name__="gelu_new",	__magic_name__=0.1,	__magic_name__=0.1,	__magic_name__=512,	__magic_name__=16,	__magic_name__=2,	__magic_name__=0.02,	__magic_name__=4,	__magic_name__="block_sparse",	__magic_name__=True,	__magic_name__=False,	__magic_name__=2,	__magic_name__=3,	)				->   Optional[int]:






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

       UpperCamelCase__							:       Dict      =   rescale_embeddings
       UpperCamelCase__							:       Tuple      =   attention_type
       UpperCamelCase__							:       Union[str, Any]      =   use_bias
       UpperCamelCase__							:       Union[str, Any]      =   block_size
       UpperCamelCase__							:       Optional[int]      =   num_random_blocks




   def    UpperCamelCase__						(					self							)				->   int:






       """simple docstring"""
       UpperCamelCase__							:       int      =   ids_tensor([self.batch_size, self.seq_length],	self.vocab_size							)

       UpperCamelCase__							:       int      =   None
       if self.use_attention_mask:
           UpperCamelCase__							:       List[str]      =   random_attention_mask([self.batch_size, self.seq_length]							)

       UpperCamelCase__							:       Optional[int]      =   None
       if self.use_token_type_ids:
           UpperCamelCase__							:       List[str]      =   ids_tensor([self.batch_size, self.seq_length],	self.type_vocab_size							)

       UpperCamelCase__							:       str      =   BigBirdConfig(
           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=__magic_name__,	initializer_range=self.initializer_range,	attention_type=self.attention_type,	block_size=self.block_size,	num_random_blocks=self.num_random_blocks,	use_bias=self.use_bias,	rescale_embeddings=self.rescale_embeddings,	)

       return config, input_ids, token_type_ids, attention_mask




   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       UpperCamelCase__							:       List[str]      =   self.prepare_config_and_inputs()
       UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__							:       Tuple      =   config_and_inputs
       UpperCamelCase__							:       Union[str, Any]      =   {
           '''input_ids''': input_ids,
           '''token_type_ids''': token_type_ids,
           '''attention_mask''': attention_mask,
       }
       return config, inputs_dict






@require_flax
class 					lowercase__ (			__lowerCamelCase    ,       unittest.TestCase       ):

   '''simple docstring'''


   a		:							Any     	=	(
       (
           FlaxBigBirdForCausalLM,
           FlaxBigBirdModel,
           FlaxBigBirdForPreTraining,
           FlaxBigBirdForMaskedLM,
           FlaxBigBirdForMultipleChoice,
           FlaxBigBirdForQuestionAnswering,
           FlaxBigBirdForSequenceClassification,
           FlaxBigBirdForTokenClassification,
       )
       if is_flax_available()
       else ()
   )

   a		:							List[str]     	=	False
   a		:							str     	=	False




   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__							:       Union[str, Any]      =   FlaxBigBirdModelTester(self							)




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Dict:






       """simple docstring"""
       super().test_from_pretrained_save_pretrained()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Optional[Any]:






       """simple docstring"""
       super().test_from_pretrained_with_no_automatic_init()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Dict:






       """simple docstring"""
       super().test_no_automatic_init()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Optional[Any]:






       """simple docstring"""
       super().test_hidden_states_output()




   @slow
   def    UpperCamelCase__						(					self							)				->   List[str]:






       """simple docstring"""
       for model_class_name in self.all_model_classes:
           UpperCamelCase__							:       List[str]      =   model_class_name.from_pretrained('''google/bigbird-roberta-base'''							)
           self.assertIsNotNone(__magic_name__							)




   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       if self.test_attn_probs:
           super().test_attention_outputs()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__     ,UpperCamelCase__							:       Dict      =   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
           with self.subTest(model_class.__name__							):
               UpperCamelCase__							:       Union[str, Any]      =   self._prepare_for_class(__magic_name__,	__magic_name__							)
               UpperCamelCase__							:       Tuple      =   model_class(__magic_name__							)

               @jax.jit
               def model_jitted(__magic_name__,	__magic_name__=None,	**__magic_name__							):
                   return model(input_ids=__magic_name__,	attention_mask=__magic_name__,	**__magic_name__							)

               with self.subTest('''JIT Enabled'''							):
                   UpperCamelCase__							:       str      =   model_jitted(**__magic_name__							).to_tuple()

               with self.subTest('''JIT Disabled'''							):
                   with jax.disable_jit():
                       UpperCamelCase__							:       Dict      =   model_jitted(**__magic_name__							).to_tuple()

               self.assertEqual(len(__magic_name__							),	len(__magic_name__							)							)
               for jitted_output, output in zip(__magic_name__,	__magic_name__							):
                   self.assertEqual(jitted_output.shape,	output.shape							)




   def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__=1E-5,	__magic_name__="outputs",	__magic_name__=None							)				->   List[Any]:






       """simple docstring"""
       # `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
       # an effort was done to return `attention_probs` (yet to be verified).
       if name.startswith('''outputs.attentions'''							):
           return
       else:
           super().check_pt_flax_outputs(__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__							)

