
import unittest

from transformers.utils.backbone_utils import (
    BackboneMixin,
    get_aligned_output_features_output_indices,
    verify_out_features_out_indices,
)





class     __UpperCamelCase (							unittest.TestCase       ):






     """simple docstring"""







     def  UpperCAmelCase__							(				self							:							str ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Any					=		['''a''', '''b''', '''c''']

       # Defaults to last layer if both are None
       __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE      :	List[str]					=		get_aligned_output_features_output_indices(_A     ,						_A     ,						_A )
       self.assertEqual(_A     ,						['''c'''] )
       self.assertEqual(_A     ,						[2] )

       # Out indices set to match out features
       __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE      :	Tuple					=		get_aligned_output_features_output_indices(['''a''', '''c''']     ,						_A     ,						_A )
       self.assertEqual(_A     ,						['''a''', '''c'''] )
       self.assertEqual(_A     ,						[0, 2] )

       # Out features set to match out indices
       __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE      :	List[Any]					=		get_aligned_output_features_output_indices(_A     ,						[0, 2]     ,						_A )
       self.assertEqual(_A     ,						['''a''', '''c'''] )
       self.assertEqual(_A     ,						[0, 2] )

       # Out features selected from negative indices
       __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE      :	Any					=		get_aligned_output_features_output_indices(_A     ,						[-3, -1]     ,						_A )
       self.assertEqual(_A     ,						['''a''', '''c'''] )
       self.assertEqual(_A     ,						[-3, -1] )



     def  UpperCAmelCase__							(				self							:							Any ):







       """simple docstring"""




       with self.assertRaises(_A ):
         verify_out_features_out_indices(['''a''', '''b''']     ,						(0, 1)     ,						_A )

       # Out features must be a list
       with self.assertRaises(_A ):
         verify_out_features_out_indices(('''a''', '''b''')     ,						(0, 1)     ,						['''a''', '''b'''] )

       # Out features must be a subset of stage names
       with self.assertRaises(_A ):
         verify_out_features_out_indices(['''a''', '''b''']     ,						(0, 1)     ,						['''a'''] )

       # Out indices must be a list or tuple
       with self.assertRaises(_A ):
         verify_out_features_out_indices(_A     ,						0     ,						['''a''', '''b'''] )

       # Out indices must be a subset of stage names
       with self.assertRaises(_A ):
         verify_out_features_out_indices(_A     ,						(0, 1)     ,						['''a'''] )

       # Out features and out indices must be the same length
       with self.assertRaises(_A ):
         verify_out_features_out_indices(['''a''', '''b''']     ,						(0,)     ,						['''a''', '''b''', '''c'''] )

       # Out features should match out indices
       with self.assertRaises(_A ):
         verify_out_features_out_indices(['''a''', '''b''']     ,						(0, 2)     ,						['''a''', '''b''', '''c'''] )

       # Out features and out indices should be in order
       with self.assertRaises(_A ):
         verify_out_features_out_indices(['''b''', '''a''']     ,						(0, 1)     ,						['''a''', '''b'''] )

       # Check passes with valid inputs
       verify_out_features_out_indices(['''a''', '''b''', '''d''']     ,						(0, 1, -1)     ,						['''a''', '''b''', '''c''', '''d'''] )




     def  UpperCAmelCase__							(				self							:							List[str] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Any					=		BackboneMixin()

       __SCREAMING_SNAKE_CASE      :	Any					=		['''a''', '''b''', '''c''']
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		['''a''', '''c''']
       __SCREAMING_SNAKE_CASE      :	int					=		[0, 2]

       # Check that the output features and indices are set correctly
       self.assertEqual(backbone.out_features     ,						['''a''', '''c'''] )
       self.assertEqual(backbone.out_indices     ,						[0, 2] )

       # Check out features and indices are updated correctly
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		['''a''', '''b''']
       self.assertEqual(backbone.out_features     ,						['''a''', '''b'''] )
       self.assertEqual(backbone.out_indices     ,						[0, 1] )

       __SCREAMING_SNAKE_CASE      :	str					=		[-3, -1]
       self.assertEqual(backbone.out_features     ,						['''a''', '''c'''] )
       self.assertEqual(backbone.out_indices     ,						[-3, -1] )

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  ):



 '''simple docstring'''




 def __init__(		self				:   List[Any]							, lowercase_				:   Union[str, Any]							, lowercase_				:   Union[str, Any]=2							, lowercase_				:   Dict=56							, lowercase_				:   str=True							, lowercase_				:   Optional[Any]=True							, lowercase_				:   List[Any]=True							, lowercase_				:   Union[str, Any]=True							, lowercase_				:   Any=99							, lowercase_				:   Optional[int]=32							, lowercase_				:   Tuple=2							, lowercase_				:   int=2							, lowercase_				:   List[str]=7							, lowercase_				:   Any="gelu_new"							, lowercase_				:   List[str]=0.1							, lowercase_				:   str=0.1							, lowercase_				:   List[Any]=512							, lowercase_				:   List[str]=16							, lowercase_				:   Optional[int]=2							, lowercase_				:   Union[str, Any]=0.02							, lowercase_				:   Union[str, Any]=4							, lowercase_				:   Union[str, Any]="block_sparse"							, lowercase_				:   Tuple=True							, lowercase_				:   Dict=False							, lowercase_				:   Dict=2							, lowercase_				:   Dict=3							, )						->	List[str]:


   """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 		__UpperCAmelCase				(		self				:   List[Any])						->	Dict:


   """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=lowercase_							, 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				:   Optional[Any])						->	Optional[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 					_UpperCAmelCase			(     lowerCAmelCase,			unittest.TestCase  ):



 '''simple docstring'''




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

 __A			       =					False
 __A			       =					False







 def 		__UpperCAmelCase				(		self				:   Optional[int])						->	Optional[Any]:


   """simple docstring"""







   _UpperCamelCase           =    FlaxBigBirdModelTester(self)







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 		__UpperCAmelCase				(		self				:   Union[str, Any])						->	List[str]:


   """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				:   Any)						->	int:


   """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				:   List[Any])						->	Union[str, Any]:


   """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				:   str)						->	Union[str, Any]:


   """simple docstring"""







   super().test_hidden_states_output()







 @slow
 def 		__UpperCAmelCase				(		self				:   Tuple)						->	Optional[int]:


   """simple docstring"""







   for model_class_name in self.all_model_classes:
     _UpperCamelCase           =    model_class_name.from_pretrained("google/bigbird-roberta-base")
     self.assertIsNotNone(lowercase_)







 def 		__UpperCAmelCase				(		self				:   Tuple)						->	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				:   Tuple)						->	Dict:


   """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(lowercase_							, lowercase_)
       _UpperCamelCase           =    model_class(lowercase_)

       @jax.jit
       def model_jitted(lowercase_				:   Dict							, lowercase_				:   List[Any]=None							, **lowercase_				:   Tuple):
         return model(input_ids=lowercase_							, attention_mask=lowercase_							, **lowercase_)

       with self.subTest("JIT Enabled"):
         _UpperCamelCase           =    model_jitted(**lowercase_).to_tuple()

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

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







 def 		__UpperCAmelCase				(		self				:   Any							, lowercase_				:   List[str]							, lowercase_				:   List[str]							, lowercase_				:   List[str]							, lowercase_				:   str=1e-5							, lowercase_				:   int="outputs"							, lowercase_				:   List[str]=None)						->	Tuple:


   """simple docstring"""







   if name.startswith("outputs.attentions"):
     return
   else:
     super().check_pt_flax_outputs(lowercase_							, lowercase_							, lowercase_							, lowercase_							, lowercase_							, lowercase_)




