






import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
 from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import ViTImageProcessor





class       __a  :







 def __init__(		self :       Optional[Any]							,       SCREAMING_SNAKE_CASE :       Tuple							,       SCREAMING_SNAKE_CASE :       str=13							,       SCREAMING_SNAKE_CASE :       Optional[int]=30							,       SCREAMING_SNAKE_CASE :       Tuple=2							,       SCREAMING_SNAKE_CASE :       int=3							,       SCREAMING_SNAKE_CASE :       Union[str, Any]=True							,       SCREAMING_SNAKE_CASE :       Optional[int]=True							,       SCREAMING_SNAKE_CASE :       Union[str, Any]=32							,       SCREAMING_SNAKE_CASE :       Tuple=5							,       SCREAMING_SNAKE_CASE :       Any=4							,       SCREAMING_SNAKE_CASE :       int=37							,       SCREAMING_SNAKE_CASE :       int="gelu"							,       SCREAMING_SNAKE_CASE :       List[Any]=0.1							,       SCREAMING_SNAKE_CASE :       str=0.1							,       SCREAMING_SNAKE_CASE :       str=10							,       SCREAMING_SNAKE_CASE :       Optional[Any]=0.0_2							,       SCREAMING_SNAKE_CASE :       int=None							,       SCREAMING_SNAKE_CASE :       Any=2							,       ):


      '''simple docstring'''






      UpperCamelCase__      :       Union[str, Any]	    =  parent
      UpperCamelCase__      :       str	    =  batch_size
      UpperCamelCase__      :       Tuple	    =  image_size
      UpperCamelCase__      :       Union[str, Any]	    =  patch_size
      UpperCamelCase__      :       Union[str, Any]	    =  num_channels
      UpperCamelCase__      :       List[Any]	    =  is_training
      UpperCamelCase__      :       int	    =  use_labels
      UpperCamelCase__      :       Any	    =  hidden_size
      UpperCamelCase__      :       str	    =  num_hidden_layers
      UpperCamelCase__      :       Optional[Any]	    =  num_attention_heads
      UpperCamelCase__      :       Union[str, Any]	    =  intermediate_size
      UpperCamelCase__      :       Tuple	    =  hidden_act
      UpperCamelCase__      :       Any	    =  hidden_dropout_prob
      UpperCamelCase__      :       Optional[Any]	    =  attention_probs_dropout_prob
      UpperCamelCase__      :       Union[str, Any]	    =  type_sequence_label_size
      UpperCamelCase__      :       int	    =  initializer_range
      UpperCamelCase__      :       Dict	    =  scope
      UpperCamelCase__      :       List[str]	    =  encoder_stride

      # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
      UpperCamelCase__      :       int	    =  (image_size // patch_size) ** 2
      UpperCamelCase__      :       str	    =  num_patches + 1







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


      '''simple docstring'''






      UpperCamelCase__      :       Optional[int]	    =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

      UpperCamelCase__      :       List[str]	    =  None
      if self.use_labels:
           UpperCamelCase__      :       Any	    =  ids_tensor([self.batch_size]							,       self.type_sequence_label_size						)

      UpperCamelCase__      :       Dict	    =  self.get_config()

      return config, pixel_values, labels







 def  __lowercase (		self :       str						):


      '''simple docstring'''






      return ViTConfig(
          image_size=self.image_size							,       patch_size=self.patch_size							,       num_channels=self.num_channels							,       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							,       is_decoder=UpperCamelCase__							,       initializer_range=self.initializer_range							,       encoder_stride=self.encoder_stride							,       )







 def  __lowercase (		self :       Dict							,       SCREAMING_SNAKE_CASE :       Optional[Any]							,       SCREAMING_SNAKE_CASE :       List[Any]							,       SCREAMING_SNAKE_CASE :       Any						):


      '''simple docstring'''






      UpperCamelCase__      :       Optional[int]	    =  ViTModel(config=UpperCamelCase__						)
      model.to(UpperCamelCase__						)
      model.eval()
      UpperCamelCase__      :       Optional[int]	    =  model(UpperCamelCase__						)
      self.parent.assertEqual(result.last_hidden_state.shape							,       (self.batch_size, self.seq_length, self.hidden_size)						)







 def  __lowercase (		self :       List[str]							,       SCREAMING_SNAKE_CASE :       List[Any]							,       SCREAMING_SNAKE_CASE :       List[Any]							,       SCREAMING_SNAKE_CASE :       str						):


      '''simple docstring'''






      UpperCamelCase__      :       int	    =  ViTForMaskedImageModeling(config=UpperCamelCase__						)
      model.to(UpperCamelCase__						)
      model.eval()
      UpperCamelCase__      :       Tuple	    =  model(UpperCamelCase__						)
      self.parent.assertEqual(
          result.reconstruction.shape							,       (self.batch_size, self.num_channels, self.image_size, self.image_size)						)

      # test greyscale images
      UpperCamelCase__      :       List[str]	    =  1
      UpperCamelCase__      :       List[str]	    =  ViTForMaskedImageModeling(UpperCamelCase__						)
      model.to(UpperCamelCase__						)
      model.eval()

      UpperCamelCase__      :       str	    =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
      UpperCamelCase__      :       List[Any]	    =  model(UpperCamelCase__						)
      self.parent.assertEqual(result.reconstruction.shape							,       (self.batch_size, 1, self.image_size, self.image_size)						)







 def  __lowercase (		self :       Tuple							,       SCREAMING_SNAKE_CASE :       Dict							,       SCREAMING_SNAKE_CASE :       List[Any]							,       SCREAMING_SNAKE_CASE :       List[Any]						):


      '''simple docstring'''






      UpperCamelCase__      :       int	    =  self.type_sequence_label_size
      UpperCamelCase__      :       str	    =  ViTForImageClassification(UpperCamelCase__						)
      model.to(UpperCamelCase__						)
      model.eval()
      UpperCamelCase__      :       Optional[int]	    =  model(UpperCamelCase__							,       labels=UpperCamelCase__						)
      self.parent.assertEqual(result.logits.shape							,       (self.batch_size, self.type_sequence_label_size)						)

      # test greyscale images
      UpperCamelCase__      :       str	    =  1
      UpperCamelCase__      :       Optional[int]	    =  ViTForImageClassification(UpperCamelCase__						)
      model.to(UpperCamelCase__						)
      model.eval()

      UpperCamelCase__      :       List[str]	    =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
      UpperCamelCase__      :       Tuple	    =  model(UpperCamelCase__						)
      self.parent.assertEqual(result.logits.shape							,       (self.batch_size, self.type_sequence_label_size)						)







 def  __lowercase (		self :       str						):


      '''simple docstring'''






      UpperCamelCase__      :       Optional[Any]	    =  self.prepare_config_and_inputs()
      (
          (
          UpperCamelCase__     
      )    ,					(
          UpperCamelCase__     
      )    ,					(
          UpperCamelCase__     
      )    ,					     
      )      :       List[str]	    =  config_and_inputs
      UpperCamelCase__      :       str	    =  {"pixel_values": pixel_values}
      return config, inputs_dict







@require_torch
class       __a  (   _snake_case	,       _snake_case	,       unittest.TestCase	):
 _lowerCAmelCase	:	Union[str, Any]				=			(
     (
         ViTModel,
         ViTForImageClassification,
         ViTForMaskedImageModeling,
     )
     if is_torch_available()
     else ()
 )
 _lowerCAmelCase	:	List[str]				=			(
     {'''feature-extraction''': ViTModel, '''image-classification''': ViTForImageClassification}
     if is_torch_available()
     else {}
 )
 _lowerCAmelCase	:	Any				=			True

 _lowerCAmelCase	:	Union[str, Any]				=			False
 _lowerCAmelCase	:	List[Any]				=			False
 _lowerCAmelCase	:	Union[str, Any]				=			False







 def  __lowercase (		self :       Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__      :       Optional[int]	    =  ViTModelTester(self						)
      UpperCamelCase__      :       Optional[int]	    =  ConfigTester(self							,       config_class=UpperCamelCase__							,       has_text_modality=UpperCamelCase__							,       hidden_size=37						)







 def  __lowercase (		self :       Tuple						):


      '''simple docstring'''






      self.config_tester.run_common_tests()







 @unittest.skip(reason="ViT does not use inputs_embeds"						)
 def  __lowercase (		self :       Optional[int]						):


      '''simple docstring'''






      pass







 def  __lowercase (		self :       Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__    ,					UpperCamelCase__      :       Any	    =  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           UpperCamelCase__      :       Tuple	    =  model_class(UpperCamelCase__						)
           self.assertIsInstance(model.get_input_embeddings()							,       (nn.Module)						)
           UpperCamelCase__      :       Optional[int]	    =  model.get_output_embeddings()
           self.assertTrue(x is None or isinstance(UpperCamelCase__							,       nn.Linear						)						)







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


      '''simple docstring'''






      UpperCamelCase__    ,					UpperCamelCase__      :       List[Any]	    =  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           UpperCamelCase__      :       Optional[int]	    =  model_class(UpperCamelCase__						)
           UpperCamelCase__      :       List[str]	    =  inspect.signature(model.forward						)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           UpperCamelCase__      :       Optional[Any]	    =  [*signature.parameters.keys()]

           UpperCamelCase__      :       Union[str, Any]	    =  ["pixel_values"]
           self.assertListEqual(arg_names[:1]							,       UpperCamelCase__						)







 def  __lowercase (		self :       str						):


      '''simple docstring'''






      UpperCamelCase__      :       int	    =  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*UpperCamelCase__						)







 def  __lowercase (		self :       Dict						):


      '''simple docstring'''






      UpperCamelCase__      :       List[Any]	    =  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_image_modeling(*UpperCamelCase__						)







 def  __lowercase (		self :       str						):


      '''simple docstring'''






      UpperCamelCase__      :       int	    =  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*UpperCamelCase__						)







 @slow
 def  __lowercase (		self :       int						):


      '''simple docstring'''






      for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           UpperCamelCase__      :       Union[str, Any]	    =  ViTModel.from_pretrained(UpperCamelCase__						)
           self.assertIsNotNone(UpperCamelCase__						)







def 		SCREAMING_SNAKE_CASE							( )			->     Tuple:
     UpperCamelCase__      :       str	    =  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"      )
     return image





@require_torch
@require_vision
class       __a  (   unittest.TestCase	):







 @cached_property
 def  __lowercase (		self :       Tuple						):


      '''simple docstring'''






      return ViTImageProcessor.from_pretrained("google/vit-base-patch16-224"						) if is_vision_available() else None







 @slow
 def  __lowercase (		self :       Dict						):


      '''simple docstring'''






      UpperCamelCase__      :       Union[str, Any]	    =  ViTForImageClassification.from_pretrained("google/vit-base-patch16-224"						).to(UpperCamelCase__						)

      UpperCamelCase__      :       List[Any]	    =  self.default_image_processor
      UpperCamelCase__      :       List[str]	    =  prepare_img()
      UpperCamelCase__      :       Tuple	    =  image_processor(images=UpperCamelCase__							,       return_tensors="pt"						).to(UpperCamelCase__						)

      # forward pass
      with torch.no_grad():
           UpperCamelCase__      :       Optional[Any]	    =  model(**UpperCamelCase__						)

      # verify the logits
      UpperCamelCase__      :       Tuple	    =  torch.Size((1, 10_00)						)
      self.assertEqual(outputs.logits.shape							,       UpperCamelCase__						)

      UpperCamelCase__      :       List[Any]	    =  torch.tensor([-0.2_7_4_4, 0.8_2_1_5, -0.0_8_3_6]						).to(UpperCamelCase__						)

      self.assertTrue(torch.allclose(outputs.logits[0, :3]							,       UpperCamelCase__							,       atol=1e-4						)						)







 @slow
 def  __lowercase (		self :       Any						):


      '''simple docstring'''






      UpperCamelCase__      :       Dict	    =  ViTModel.from_pretrained("facebook/dino-vits8"						).to(UpperCamelCase__						)

      UpperCamelCase__      :       Optional[int]	    =  ViTImageProcessor.from_pretrained("facebook/dino-vits8"							,       size=4_80						)
      UpperCamelCase__      :       Dict	    =  prepare_img()
      UpperCamelCase__      :       str	    =  image_processor(images=UpperCamelCase__							,       return_tensors="pt"						)
      UpperCamelCase__      :       List[str]	    =  inputs.pixel_values.to(UpperCamelCase__						)

      # forward pass
      with torch.no_grad():
           UpperCamelCase__      :       List[Any]	    =  model(UpperCamelCase__							,       interpolate_pos_encoding=UpperCamelCase__						)

      # verify the logits
      UpperCamelCase__      :       Dict	    =  torch.Size((1, 36_01, 3_84)						)
      self.assertEqual(outputs.last_hidden_state.shape							,       UpperCamelCase__						)

      UpperCamelCase__      :       int	    =  torch.tensor(
          [[4.2_3_4_0, 4.3_9_0_6, -6.6_6_9_2], [4.5_4_6_3, 1.8_9_2_8, -6.7_2_5_7], [4.4_4_2_9, 0.8_4_9_6, -5.8_5_8_5]]						).to(UpperCamelCase__						)

      self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]							,       UpperCamelCase__							,       atol=1e-4						)						)







 @slow
 @require_accelerate
 @require_torch_gpu
 def  __lowercase (		self :       Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__      :       List[Any]	    =  ViTModel.from_pretrained("facebook/dino-vits8"							,       torch_dtype=torch.floataa							,       device_map="auto"						)
      UpperCamelCase__      :       Optional[Any]	    =  self.default_image_processor

      UpperCamelCase__      :       int	    =  prepare_img()
      UpperCamelCase__      :       Optional[Any]	    =  image_processor(images=UpperCamelCase__							,       return_tensors="pt"						)
      UpperCamelCase__      :       str	    =  inputs.pixel_values.to(UpperCamelCase__						)

      # forward pass to make sure inference works in fp16
      with torch.no_grad():
           UpperCamelCase__      :       int	    =  model(UpperCamelCase__						)







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







 def __init__(		self :       int							,       SCREAMING_SNAKE_CASE :       Optional[Any]							,       SCREAMING_SNAKE_CASE :       Dict=2							,       SCREAMING_SNAKE_CASE :       List[Any]=56							,       SCREAMING_SNAKE_CASE :       List[str]=True							,       SCREAMING_SNAKE_CASE :       str=True							,       SCREAMING_SNAKE_CASE :       Any=True							,       SCREAMING_SNAKE_CASE :       int=True							,       SCREAMING_SNAKE_CASE :       List[str]=99							,       SCREAMING_SNAKE_CASE :       str=32							,       SCREAMING_SNAKE_CASE :       Optional[int]=2							,       SCREAMING_SNAKE_CASE :       Union[str, Any]=2							,       SCREAMING_SNAKE_CASE :       Dict=7							,       SCREAMING_SNAKE_CASE :       List[Any]="gelu_new"							,       SCREAMING_SNAKE_CASE :       List[Any]=0.1							,       SCREAMING_SNAKE_CASE :       int=0.1							,       SCREAMING_SNAKE_CASE :       Any=5_12							,       SCREAMING_SNAKE_CASE :       Dict=16							,       SCREAMING_SNAKE_CASE :       Union[str, Any]=2							,       SCREAMING_SNAKE_CASE :       Any=0.0_2							,       SCREAMING_SNAKE_CASE :       Optional[Any]=4							,       SCREAMING_SNAKE_CASE :       int="block_sparse"							,       SCREAMING_SNAKE_CASE :       int=True							,       SCREAMING_SNAKE_CASE :       Optional[Any]=False							,       SCREAMING_SNAKE_CASE :       Tuple=2							,       SCREAMING_SNAKE_CASE :       Dict=3							,       ):


      '''simple docstring'''






      UpperCamelCase__      :       List[str]	    =  parent
      UpperCamelCase__      :       Union[str, Any]	    =  batch_size
      UpperCamelCase__      :       Union[str, Any]	    =  seq_length
      UpperCamelCase__      :       Dict	    =  is_training
      UpperCamelCase__      :       Optional[int]	    =  use_attention_mask
      UpperCamelCase__      :       List[str]	    =  use_token_type_ids
      UpperCamelCase__      :       Dict	    =  use_labels
      UpperCamelCase__      :       Optional[int]	    =  vocab_size
      UpperCamelCase__      :       List[Any]	    =  hidden_size
      UpperCamelCase__      :       List[str]	    =  num_hidden_layers
      UpperCamelCase__      :       List[str]	    =  num_attention_heads
      UpperCamelCase__      :       int	    =  intermediate_size
      UpperCamelCase__      :       str	    =  hidden_act
      UpperCamelCase__      :       Tuple	    =  hidden_dropout_prob
      UpperCamelCase__      :       Any	    =  attention_probs_dropout_prob
      UpperCamelCase__      :       str	    =  max_position_embeddings
      UpperCamelCase__      :       Tuple	    =  type_vocab_size
      UpperCamelCase__      :       Dict	    =  type_sequence_label_size
      UpperCamelCase__      :       Optional[Any]	    =  initializer_range
      UpperCamelCase__      :       Any	    =  num_choices

      UpperCamelCase__      :       Dict	    =  rescale_embeddings
      UpperCamelCase__      :       Union[str, Any]	    =  attention_type
      UpperCamelCase__      :       int	    =  use_bias
      UpperCamelCase__      :       List[Any]	    =  block_size
      UpperCamelCase__      :       Union[str, Any]	    =  num_random_blocks







 def  __lowercase (		self :       List[str]						):


      '''simple docstring'''






      UpperCamelCase__      :       Dict	    =  ids_tensor([self.batch_size, self.seq_length]							,       self.vocab_size						)

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

      UpperCamelCase__      :       Any	    =  None
      if self.use_token_type_ids:
           UpperCamelCase__      :       int	    =  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  __lowercase (		self :       Any						):


      '''simple docstring'''






      UpperCamelCase__      :       int	    =  self.prepare_config_and_inputs()
      UpperCamelCase__    ,					UpperCamelCase__    ,					UpperCamelCase__    ,					UpperCamelCase__      :       Tuple	    =  config_and_inputs
      UpperCamelCase__      :       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	):
 _lowerCAmelCase	:	Optional[Any]				=			(
     (
         FlaxBigBirdForCausalLM,
         FlaxBigBirdModel,
         FlaxBigBirdForPreTraining,
         FlaxBigBirdForMaskedLM,
         FlaxBigBirdForMultipleChoice,
         FlaxBigBirdForQuestionAnswering,
         FlaxBigBirdForSequenceClassification,
         FlaxBigBirdForTokenClassification,
     )
     if is_flax_available()
     else ()
 )

 _lowerCAmelCase	:	Dict				=			False
 _lowerCAmelCase	:	Optional[int]				=			False







 def  __lowercase (		self :       Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__      :       Any	    =  FlaxBigBirdModelTester(self						)







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


      '''simple docstring'''






      super().test_no_automatic_init()







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


      '''simple docstring'''






      super().test_hidden_states_output()







 @slow
 def  __lowercase (		self :       str						):


      '''simple docstring'''






      for model_class_name in self.all_model_classes:
           UpperCamelCase__      :       Any	    =  model_class_name.from_pretrained("google/bigbird-roberta-base"						)
           self.assertIsNotNone(SCREAMING_SNAKE_CASE						)







 def  __lowercase (		self :       Optional[int]						):


      '''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  __lowercase (		self :       Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__    ,					UpperCamelCase__      :       Tuple	    =  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           with self.subTest(model_class.__name__						):
                UpperCamelCase__      :       Dict	    =  self._prepare_for_class(SCREAMING_SNAKE_CASE							,       SCREAMING_SNAKE_CASE						)
                UpperCamelCase__      :       Union[str, Any]	    =  model_class(SCREAMING_SNAKE_CASE						)

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

                with self.subTest("JIT Enabled"						):
                     UpperCamelCase__      :       Tuple	    =  model_jitted(**SCREAMING_SNAKE_CASE						).to_tuple()

                with self.subTest("JIT Disabled"						):
                     with jax.disable_jit():
                          UpperCamelCase__      :       List[Any]	    =  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  __lowercase (		self :       int							,       SCREAMING_SNAKE_CASE :       int							,       SCREAMING_SNAKE_CASE :       int							,       SCREAMING_SNAKE_CASE :       Union[str, Any]							,       SCREAMING_SNAKE_CASE :       Any=1e-5							,       SCREAMING_SNAKE_CASE :       Tuple="outputs"							,       SCREAMING_SNAKE_CASE :       Optional[Any]=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						)