




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     __a(	unittest.TestCase  ):






    """simple docstring"""





    def __init__( self  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE=2  ,_SCREAMING_SNAKE_CASE=56  ,_SCREAMING_SNAKE_CASE=True  ,_SCREAMING_SNAKE_CASE=True  ,_SCREAMING_SNAKE_CASE=True  ,_SCREAMING_SNAKE_CASE=True  ,_SCREAMING_SNAKE_CASE=99  ,_SCREAMING_SNAKE_CASE=32  ,_SCREAMING_SNAKE_CASE=2  ,_SCREAMING_SNAKE_CASE=2  ,_SCREAMING_SNAKE_CASE=7  ,_SCREAMING_SNAKE_CASE="gelu_new"  ,_SCREAMING_SNAKE_CASE=0.1  ,_SCREAMING_SNAKE_CASE=0.1  ,_SCREAMING_SNAKE_CASE=512  ,_SCREAMING_SNAKE_CASE=16  ,_SCREAMING_SNAKE_CASE=2  ,_SCREAMING_SNAKE_CASE=0.02  ,_SCREAMING_SNAKE_CASE=4  ,_SCREAMING_SNAKE_CASE="block_sparse"  ,_SCREAMING_SNAKE_CASE=True  ,_SCREAMING_SNAKE_CASE=False  ,_SCREAMING_SNAKE_CASE=2  ,_SCREAMING_SNAKE_CASE=3  ,)					->						Tuple:
         UpperCAmelCase_     :				int  =     parent
         UpperCAmelCase_     :				str  =     batch_size
         UpperCAmelCase_     :				Tuple  =     seq_length
         UpperCAmelCase_     :				Optional[Any]  =     is_training
         UpperCAmelCase_     :				Union[str, Any]  =     use_attention_mask
         UpperCAmelCase_     :				Union[str, Any]  =     use_token_type_ids
         UpperCAmelCase_     :				List[str]  =     use_labels
         UpperCAmelCase_     :				List[str]  =     vocab_size
         UpperCAmelCase_     :				Tuple  =     hidden_size
         UpperCAmelCase_     :				Optional[Any]  =     num_hidden_layers
         UpperCAmelCase_     :				List[Any]  =     num_attention_heads
         UpperCAmelCase_     :				str  =     intermediate_size
         UpperCAmelCase_     :				Dict  =     hidden_act
         UpperCAmelCase_     :				Optional[int]  =     hidden_dropout_prob
         UpperCAmelCase_     :				Tuple  =     attention_probs_dropout_prob
         UpperCAmelCase_     :				str  =     max_position_embeddings
         UpperCAmelCase_     :				List[str]  =     type_vocab_size
         UpperCAmelCase_     :				Any  =     type_sequence_label_size
         UpperCAmelCase_     :				List[Any]  =     initializer_range
         UpperCAmelCase_     :				List[str]  =     num_choices

         UpperCAmelCase_     :				int  =     rescale_embeddings
         UpperCAmelCase_     :				List[Any]  =     attention_type
         UpperCAmelCase_     :				Dict  =     use_bias
         UpperCAmelCase_     :				Dict  =     block_size
         UpperCAmelCase_     :				List[str]  =     num_random_blocks





    def 				a__ ( self    )					->						Union[str, Any]:
         UpperCAmelCase_     :				int  =     ids_tensor([self.batch_size, self.seq_length]  ,self.vocab_size    )

         UpperCAmelCase_     :				Tuple  =     None
         if self.use_attention_mask:
              UpperCAmelCase_     :				Union[str, Any]  =     random_attention_mask([self.batch_size, self.seq_length]    )

         UpperCAmelCase_     :				Optional[Any]  =     None
         if self.use_token_type_ids:
              UpperCAmelCase_     :				Dict  =     ids_tensor([self.batch_size, self.seq_length]  ,self.type_vocab_size    )

         UpperCAmelCase_     :				List[Any]  =     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=_SCREAMING_SNAKE_CASE  ,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 				a__ ( self    )					->						Optional[Any]:
         UpperCAmelCase_     :				Dict  =     self.prepare_config_and_inputs()
         UpperCAmelCase_,       UpperCAmelCase_,       UpperCAmelCase_,       UpperCAmelCase_     :				Optional[Any]  =     config_and_inputs
         UpperCAmelCase_     :				Optional[int]  =     {
             '''input_ids''': input_ids,
             '''token_type_ids''': token_type_ids,
             '''attention_mask''': attention_mask,
         }
         return config, inputs_dict


@require_flax
class     __a(	_a							,		unittest.TestCase  ):






    """simple docstring"""





    lowerCAmelCase         = (
        (
            FlaxBigBirdForCausalLM,
            FlaxBigBirdModel,
            FlaxBigBirdForPreTraining,
            FlaxBigBirdForMaskedLM,
            FlaxBigBirdForMultipleChoice,
            FlaxBigBirdForQuestionAnswering,
            FlaxBigBirdForSequenceClassification,
            FlaxBigBirdForTokenClassification,
        )
        if is_flax_available()
        else ()
    )

    lowerCAmelCase         = False
    lowerCAmelCase         = False





    def 				a__ ( self    )					->						int:
         UpperCAmelCase_     :				int  =     FlaxBigBirdModelTester(self    )





    @slow
    # copied from `test_modeling_flax_common` because it takes much longer than other models
    def 				a__ ( self    )					->						Optional[int]:
         super().test_from_pretrained_save_pretrained()





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





    @slow
    # copied from `test_modeling_flax_common` because it takes much longer than other models
    def 				a__ ( self    )					->						Tuple:
         super().test_no_automatic_init()





    @slow
    # copied from `test_modeling_flax_common` because it takes much longer than other models
    def 				a__ ( self    )					->						int:
         super().test_hidden_states_output()





    @slow
    def 				a__ ( self    )					->						str:
         for model_class_name in self.all_model_classes:
              UpperCAmelCase_     :				str  =     model_class_name.from_pretrained('''google/bigbird-roberta-base'''    )
              self.assertIsNotNone(_SCREAMING_SNAKE_CASE    )





    def 				a__ ( self    )					->						Any:
         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 				a__ ( self    )					->						List[str]:
         UpperCAmelCase_,       UpperCAmelCase_     :				Optional[Any]  =     self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__    ):
                   UpperCAmelCase_     :				Optional[int]  =     self._prepare_for_class(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )
                   UpperCAmelCase_     :				Optional[Any]  =     model_class(_SCREAMING_SNAKE_CASE    )

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

                   with self.subTest('''JIT Enabled'''    ):
                        UpperCAmelCase_     :				Optional[int]  =     model_jitted(**_SCREAMING_SNAKE_CASE    ).to_tuple()

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

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





    def 				a__ ( self  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE=1e-5  ,_SCREAMING_SNAKE_CASE="outputs"  ,_SCREAMING_SNAKE_CASE=None    )					->						Any:
         # `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(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )





def lowerCamelCase__  (     ):

     '''simple docstring'''


     UpperCAmelCase_     :				Dict  =     0
     for i in range(1					,		1001 ):
          total += i**i
     return str(_lowercase )[-10:]


if __name__ == "__main__":
     print(solution())