


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

										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		rescale_embeddings
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		attention_type
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		use_bias
										__SCREAMING_SNAKE_CASE						:     Any					=		block_size
										__SCREAMING_SNAKE_CASE						:     Dict					=		num_random_blocks


				def 		__magic_name__(       self			:Tuple	)		->			List[str]:
										__SCREAMING_SNAKE_CASE						:     int					=		ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size	)

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

										__SCREAMING_SNAKE_CASE						:     Tuple					=		None
										if self.use_token_type_ids:
																__SCREAMING_SNAKE_CASE						:     Optional[int]					=		ids_tensor([self.batch_size, self.seq_length]			,			self.type_vocab_size	)

										__SCREAMING_SNAKE_CASE						:     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=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 		__magic_name__(       self			:List[Any]	)		->			Any:
										__SCREAMING_SNAKE_CASE						:     int					=		self.prepare_config_and_inputs()
										__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     int					=		config_and_inputs
										__SCREAMING_SNAKE_CASE						:     str					=		{
										    '''input_ids''': input_ids,
										    '''token_type_ids''': token_type_ids,
										    '''attention_mask''': attention_mask,
										}
										return config, inputs_dict




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






				'''simple docstring'''



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

				SCREAMING_SNAKE_CASE__				:			Optional[int]						=		False
				SCREAMING_SNAKE_CASE__				:			Optional[Any]						=		False


				def 		__magic_name__(       self			:Union[str, Any]	)		->			Dict:
										__SCREAMING_SNAKE_CASE						:     str					=		FlaxBigBirdModelTester(self	)


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


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


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


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


				@slow
				def 		__magic_name__(       self			:Optional[Any]	)		->			Tuple:
										for model_class_name in self.all_model_classes:
																__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		model_class_name.from_pretrained('''google/bigbird-roberta-base'''	)
																self.assertIsNotNone(lowerCAmelCase__	)


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

										for model_class in self.all_model_classes:
																with self.subTest(model_class.__name__	):
																						__SCREAMING_SNAKE_CASE						:     Any					=		self._prepare_for_class(lowerCAmelCase__			,			lowerCAmelCase__	)
																						__SCREAMING_SNAKE_CASE						:     str					=		model_class(lowerCAmelCase__	)

																						@jax.jit
																						def model_jitted(lowerCAmelCase__			:Union[str, Any]			,			lowerCAmelCase__			:Optional[int]=None			,			**lowerCAmelCase__			:Union[str, Any]	):
																												return model(input_ids=lowerCAmelCase__			,			attention_mask=lowerCAmelCase__			,			**lowerCAmelCase__	)

																						with self.subTest('''JIT Enabled'''	):
																												__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		model_jitted(**lowerCAmelCase__	).to_tuple()

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








from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_speech_available, is_torch_available


lowerCamelCase	:							int				={
    '''configuration_audio_spectrogram_transformer''': [
        '''AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''ASTConfig''',
    ]
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase	:							Union[str, Any]				=[
     '''AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''ASTForAudioClassification''',
     '''ASTModel''',
     '''ASTPreTrainedModel''',
 ]

try:
 if not is_speech_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase	:							Optional[int]				=['''ASTFeatureExtractor''']

if TYPE_CHECKING:
 from .configuration_audio_spectrogram_transformer import (
     AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
     ASTConfig,
 )

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_audio_spectrogram_transformer import (
      AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
      ASTForAudioClassification,
      ASTModel,
      ASTPreTrainedModel,
  )

 try:
  if not is_speech_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .feature_extraction_audio_spectrogram_transformer import ASTFeatureExtractor


else:
 import sys

 lowerCamelCase	:							Optional[int]				=_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)