






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      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Any   ,						_SCREAMING_SNAKE_CASE      :    str=2   ,						_SCREAMING_SNAKE_CASE      :    int=56   ,						_SCREAMING_SNAKE_CASE      :    Tuple=True   ,						_SCREAMING_SNAKE_CASE      :    int=True   ,						_SCREAMING_SNAKE_CASE      :    str=True   ,						_SCREAMING_SNAKE_CASE      :    Dict=True   ,						_SCREAMING_SNAKE_CASE      :    List[Any]=99   ,						_SCREAMING_SNAKE_CASE      :    int=32   ,						_SCREAMING_SNAKE_CASE      :    List[str]=2   ,						_SCREAMING_SNAKE_CASE      :    str=2   ,						_SCREAMING_SNAKE_CASE      :    Any=7   ,						_SCREAMING_SNAKE_CASE      :    Any="gelu_new"   ,						_SCREAMING_SNAKE_CASE      :    List[str]=0.1   ,						_SCREAMING_SNAKE_CASE      :    Any=0.1   ,						_SCREAMING_SNAKE_CASE      :    Any=512   ,						_SCREAMING_SNAKE_CASE      :    List[Any]=16   ,						_SCREAMING_SNAKE_CASE      :    List[Any]=2   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]=0.0_2   ,						_SCREAMING_SNAKE_CASE      :    Any=4   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]="block_sparse"   ,						_SCREAMING_SNAKE_CASE      :    Dict=True   ,						_SCREAMING_SNAKE_CASE      :    Tuple=False   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=2   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=3   ,						):






           """simple docstring"""




           UpperCamelCase					       =  parent
           UpperCamelCase					       =  batch_size
           UpperCamelCase					       =  seq_length
           UpperCamelCase					       =  is_training
           UpperCamelCase					       =  use_attention_mask
           UpperCamelCase					       =  use_token_type_ids
           UpperCamelCase					       =  use_labels
           UpperCamelCase					       =  vocab_size
           UpperCamelCase					       =  hidden_size
           UpperCamelCase					       =  num_hidden_layers
           UpperCamelCase					       =  num_attention_heads
           UpperCamelCase					       =  intermediate_size
           UpperCamelCase					       =  hidden_act
           UpperCamelCase					       =  hidden_dropout_prob
           UpperCamelCase					       =  attention_probs_dropout_prob
           UpperCamelCase					       =  max_position_embeddings
           UpperCamelCase					       =  type_vocab_size
           UpperCamelCase					       =  type_sequence_label_size
           UpperCamelCase					       =  initializer_range
           UpperCamelCase					       =  num_choices

           UpperCamelCase					       =  rescale_embeddings
           UpperCamelCase					       =  attention_type
           UpperCamelCase					       =  use_bias
           UpperCamelCase					       =  block_size
           UpperCamelCase					       =  num_random_blocks







     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  ids_tensor([self.batch_size, self.seq_length]   ,						self.vocab_size )

           UpperCamelCase					       =  None
           if self.use_attention_mask:
                 UpperCamelCase					       =  random_attention_mask([self.batch_size, self.seq_length] )

           UpperCamelCase					       =  None
           if self.use_token_type_ids:
                 UpperCamelCase					       =  ids_tensor([self.batch_size, self.seq_length]   ,						self.type_vocab_size )

           UpperCamelCase					       =  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   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  self.prepare_config_and_inputs()
           UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  config_and_inputs
           UpperCamelCase					       =  {
               'input_ids': input_ids,
               'token_type_ids': token_type_ids,
               'attention_mask': attention_mask,
           }
           return config, inputs_dict





@require_flax
class 					A__     (     __snake_case   ,    unittest.TestCase ):

     '''simple docstring'''






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

     snake_case__					=					False
     snake_case__					=					False







     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  FlaxBigBirdModelTester(self )







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






           """simple docstring"""




           super().test_from_pretrained_save_pretrained()







     @slow
     # copied from `test_modeling_flax_common` because it takes much longer than other models
     def   _SCREAMING_SNAKE_CASE     ( self      :    List[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   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           super().test_no_automatic_init()







     @slow
     # copied from `test_modeling_flax_common` because it takes much longer than other models
     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str] ):






           """simple docstring"""




           super().test_hidden_states_output()







     @slow
     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           for model_class_name in self.all_model_classes:
                 UpperCamelCase					       =  model_class_name.from_pretrained('google/bigbird-roberta-base' )
                 self.assertIsNotNone(_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """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   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           UpperCamelCase       ,  UpperCamelCase					       =  self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 with self.subTest(model_class.__name__ ):
                       UpperCamelCase					       =  self._prepare_for_class(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                       UpperCamelCase					       =  model_class(_SCREAMING_SNAKE_CASE )

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

                       with self.subTest('JIT Enabled' ):
                             UpperCamelCase					       =  model_jitted(**_SCREAMING_SNAKE_CASE ).to_tuple()

                       with self.subTest('JIT Disabled' ):
                             with jax.disable_jit():
                                   UpperCamelCase					       =  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   _SCREAMING_SNAKE_CASE     ( self      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]=1E-5   ,						_SCREAMING_SNAKE_CASE      :    Tuple="outputs"   ,						_SCREAMING_SNAKE_CASE      :    Dict=None ):






           """simple docstring"""




           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 )











__magic_name__   :  str        =							8.314_4598
def 		lowercase__  (				_UpperCamelCase       ,							_UpperCamelCase) ->			float:
      """simple docstring"""

      if temperature < 0:
            raise Exception('Temperature cannot be less than 0 K')
      if molar_mass <= 0:
            raise Exception('Molar mass cannot be less than or equal to 0 kg/mol')
      else:
            return (3 * UNIVERSAL_GAS_CONSTANT * temperature / molar_mass) ** 0.5


if __name__ == "__main__":
       import doctest

       # run doctest
       doctest.testmod()

       # example
       __magic_name__   :  List[Any]        =							300
       __magic_name__   :  Union[str, Any]        =							28
       __magic_name__   :  str        =							rms_speed_of_molecule(temperature, molar_mass)
       print(F'Vrms of Nitrogen gas at 300 K is {vrms} m/s')



