"""simple docstring"""


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

				def __init__( self      ,					_lowerCamelCase      ,					_lowerCamelCase=2      ,					_lowerCamelCase=56      ,					_lowerCamelCase=True      ,					_lowerCamelCase=True      ,					_lowerCamelCase=True      ,					_lowerCamelCase=True      ,					_lowerCamelCase=99      ,					_lowerCamelCase=32      ,					_lowerCamelCase=2      ,					_lowerCamelCase=2      ,					_lowerCamelCase=7      ,					_lowerCamelCase="gelu_new"      ,					_lowerCamelCase=0.1      ,					_lowerCamelCase=0.1      ,					_lowerCamelCase=512      ,					_lowerCamelCase=16      ,					_lowerCamelCase=2      ,					_lowerCamelCase=0.0_2      ,					_lowerCamelCase=4      ,					_lowerCamelCase="block_sparse"      ,					_lowerCamelCase=True      ,					_lowerCamelCase=False      ,					_lowerCamelCase=2      ,					_lowerCamelCase=3      ,					):
										lowerCamelCase__   =parent
										lowerCamelCase__   =batch_size
										lowerCamelCase__   =seq_length
										lowerCamelCase__   =is_training
										lowerCamelCase__   =use_attention_mask
										lowerCamelCase__   =use_token_type_ids
										lowerCamelCase__   =use_labels
										lowerCamelCase__   =vocab_size
										lowerCamelCase__   =hidden_size
										lowerCamelCase__   =num_hidden_layers
										lowerCamelCase__   =num_attention_heads
										lowerCamelCase__   =intermediate_size
										lowerCamelCase__   =hidden_act
										lowerCamelCase__   =hidden_dropout_prob
										lowerCamelCase__   =attention_probs_dropout_prob
										lowerCamelCase__   =max_position_embeddings
										lowerCamelCase__   =type_vocab_size
										lowerCamelCase__   =type_sequence_label_size
										lowerCamelCase__   =initializer_range
										lowerCamelCase__   =num_choices

										lowerCamelCase__   =rescale_embeddings
										lowerCamelCase__   =attention_type
										lowerCamelCase__   =use_bias
										lowerCamelCase__   =block_size
										lowerCamelCase__   =num_random_blocks

				def 						_a      ( self ):
										lowerCamelCase__   =ids_tensor([self.batch_size, self.seq_length]      ,					self.vocab_size )

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

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

										lowerCamelCase__   =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=_lowerCamelCase      ,					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 ):
										lowerCamelCase__   =self.prepare_config_and_inputs()
										lowerCamelCase__   , lowerCamelCase__   , lowerCamelCase__   , lowerCamelCase__   =config_and_inputs
										lowerCamelCase__   ={
										    "input_ids": input_ids,
										    "token_type_ids": token_type_ids,
										    "attention_mask": attention_mask,
										}
										return config, inputs_dict






@require_flax
class   __UpperCAmelCase		(      __lowerCAmelCase , unittest.TestCase      ):
				A__ :  List[str]    				=			(
				    (
				        FlaxBigBirdForCausalLM,
				        FlaxBigBirdModel,
				        FlaxBigBirdForPreTraining,
				        FlaxBigBirdForMaskedLM,
				        FlaxBigBirdForMultipleChoice,
				        FlaxBigBirdForQuestionAnswering,
				        FlaxBigBirdForSequenceClassification,
				        FlaxBigBirdForTokenClassification,
				    )
				    if is_flax_available()
				    else ()
				)

				A__ :  List[str]    				=			False
				A__ :  Optional[Any]    				=			False

				def 						_a      ( self ):
										lowerCamelCase__   =FlaxBigBirdModelTester(self )

				@slow
				# copied from `test_modeling_flax_common` because it takes much longer than other models
				def 						_a      ( self ):
										super().test_from_pretrained_save_pretrained()

				@slow
				# copied from `test_modeling_flax_common` because it takes much longer than other models
				def 						_a      ( self ):
										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 ):
										super().test_no_automatic_init()

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

				@slow
				def 						_a      ( self ):
										for model_class_name in self.all_model_classes:
																lowerCamelCase__   =model_class_name.from_pretrained("google/bigbird-roberta-base" )
																self.assertIsNotNone(_lowerCamelCase )

				def 						_a      ( self ):
										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 ):
										lowerCamelCase__   , lowerCamelCase__   =self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																with self.subTest(model_class.__name__ ):
																						lowerCamelCase__   =self._prepare_for_class(_lowerCamelCase      ,					_lowerCamelCase )
																						lowerCamelCase__   =model_class(_lowerCamelCase )

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

																						with self.subTest("JIT Enabled" ):
																												lowerCamelCase__   =model_jitted(**_lowerCamelCase ).to_tuple()

																						with self.subTest("JIT Disabled" ):
																												with jax.disable_jit():
																																		lowerCamelCase__   =model_jitted(**_lowerCamelCase ).to_tuple()

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






				def 						_a      ( self      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase=1E-5      ,					_lowerCamelCase="outputs"      ,					_lowerCamelCase=None ):
										# `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(_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase )



"""simple docstring"""


import os
from distutils.util import strtobool





def      lowerCamelCase_					(   __lowerCAmelCase  ,       __lowerCAmelCase							)				->   Union[str, Any]:




						'''simple docstring'''



						for e in env_keys:
												lowerCamelCase__   =int(os.environ.get(__lowerCAmelCase  ,       -1							)							)
												if val >= 0:
																		return val
						return default





def      lowerCamelCase_					(   __lowerCAmelCase  ,       __lowerCAmelCase=False							)				->   List[str]:




						'''simple docstring'''



						lowerCamelCase__   =os.environ.get(__lowerCAmelCase  ,       str(__lowerCAmelCase							)							)
						return strtobool(__lowerCAmelCase							) == 1  # As its name indicates `strtobool` actually returns an int...





def      lowerCamelCase_					(   __lowerCAmelCase  ,       __lowerCAmelCase="no"							)				->   Union[str, Any]:




						'''simple docstring'''



						lowerCamelCase__   =os.environ.get(__lowerCAmelCase  ,       str(__lowerCAmelCase							)							)
						return value


