





import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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


if is_torch_available():
    import torch

    from transformers import (
        MODEL_FOR_PRETRAINING_MAPPING,
        MegatronBertForCausalLM,
        MegatronBertForMaskedLM,
        MegatronBertForMultipleChoice,
        MegatronBertForNextSentencePrediction,
        MegatronBertForPreTraining,
        MegatronBertForQuestionAnswering,
        MegatronBertForSequenceClassification,
        MegatronBertForTokenClassification,
        MegatronBertModel,
    )





class  A__    :




   """simple docstring"""



   def __init__(					self					,  lowercase					,  lowercase=13					,  lowercase=7					,  lowercase=True					,  lowercase=True					,  lowercase=True					,  lowercase=True					,  lowercase=99					,  lowercase=64					,  lowercase=32					,  lowercase=5					,  lowercase=4					,  lowercase=37					,  lowercase="gelu"					,  lowercase=0.1					,  lowercase=0.1					,  lowercase=512					,  lowercase=16					,  lowercase=2					,  lowercase=0.02					,  lowercase=3					,  lowercase=4					,  lowercase=None					,  )			->						Tuple:







          '''simple docstring'''

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

   def 							__lowercase		(					self)			->						Optional[int]:







          '''simple docstring'''

          a__						:	Any =				ids_tensor([self.batch_size, self.seq_length]					,  self.vocab_size)

          a__						:	Dict =				None
          if self.use_input_mask:
                 a__						:	Union[str, Any] =				random_attention_mask([self.batch_size, self.seq_length])

          a__						:	Dict =				None
          if self.use_token_type_ids:
                 a__						:	Union[str, Any] =				ids_tensor([self.batch_size, self.seq_length]					,  self.type_vocab_size)

          a__						:	Union[str, Any] =				None
          a__						:	Union[str, Any] =				None
          a__						:	str =				None
          if self.use_labels:
                 a__						:	Optional[int] =				ids_tensor([self.batch_size]					,  self.type_sequence_label_size)
                 a__						:	Union[str, Any] =				ids_tensor([self.batch_size, self.seq_length]					,  self.num_labels)
                 a__						:	Tuple =				ids_tensor([self.batch_size]					,  self.num_choices)

          a__						:	List[str] =				self.get_config()

          return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

   def 							__lowercase		(					self)			->						int:







          '''simple docstring'''

          return MegatronBertConfig(
              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					,  embedding_size=self.embedding_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					,  )

   def 							__lowercase		(					self					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase)			->						Optional[int]:







          '''simple docstring'''

          a__						:	Union[str, Any] =				MegatronBertModel(config=lowercase)
          model.to(lowercase)
          model.eval()
          a__						:	List[str] =				model(lowercase					,  attention_mask=lowercase					,  token_type_ids=lowercase)
          a__						:	Any =				model(lowercase					,  token_type_ids=lowercase)
          a__						:	List[str] =				model(lowercase)

          self.parent.assertEqual(result.last_hidden_state.shape					,  (self.batch_size, self.seq_length, self.hidden_size))
          self.parent.assertEqual(result.pooler_output.shape					,  (self.batch_size, self.hidden_size))

   def 							__lowercase		(					self					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase)			->						Optional[int]:







          '''simple docstring'''

          a__						:	int =				MegatronBertForMaskedLM(config=lowercase)
          model.to(lowercase)
          model.eval()
          a__						:	List[Any] =				model(lowercase					,  attention_mask=lowercase					,  token_type_ids=lowercase					,  labels=lowercase)
          self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.seq_length, self.vocab_size))

   def 							__lowercase		(					self					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase)			->						List[Any]:







          '''simple docstring'''

          a__						:	Optional[Any] =				MegatronBertForCausalLM(config=lowercase)
          model.to(lowercase)
          model.eval()
          a__						:	Dict =				model(lowercase					,  attention_mask=lowercase					,  token_type_ids=lowercase					,  labels=lowercase)
          self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.seq_length, self.vocab_size))

   def 							__lowercase		(					self					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase)			->						Any:







          '''simple docstring'''

          a__						:	Optional[int] =				MegatronBertForNextSentencePrediction(config=lowercase)
          model.to(lowercase)
          model.eval()
          a__						:	Any =				model(
              lowercase					,  attention_mask=lowercase					,  token_type_ids=lowercase					,  labels=lowercase					,  )
          self.parent.assertEqual(result.logits.shape					,  (self.batch_size, 2))

   def 							__lowercase		(					self					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase)			->						Dict:







          '''simple docstring'''

          a__						:	Tuple =				MegatronBertForPreTraining(config=lowercase)
          model.to(lowercase)
          model.eval()
          a__						:	Optional[Any] =				model(
              lowercase					,  attention_mask=lowercase					,  token_type_ids=lowercase					,  labels=lowercase					,  next_sentence_label=lowercase					,  )
          self.parent.assertEqual(result.prediction_logits.shape					,  (self.batch_size, self.seq_length, self.vocab_size))
          self.parent.assertEqual(result.seq_relationship_logits.shape					,  (self.batch_size, 2))

   def 							__lowercase		(					self					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase)			->						Optional[int]:







          '''simple docstring'''

          a__						:	List[str] =				MegatronBertForQuestionAnswering(config=lowercase)
          model.to(lowercase)
          model.eval()
          a__						:	str =				model(
              lowercase					,  attention_mask=lowercase					,  token_type_ids=lowercase					,  start_positions=lowercase					,  end_positions=lowercase					,  )
          self.parent.assertEqual(result.start_logits.shape					,  (self.batch_size, self.seq_length))
          self.parent.assertEqual(result.end_logits.shape					,  (self.batch_size, self.seq_length))

   def 							__lowercase		(					self					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase)			->						List[str]:







          '''simple docstring'''

          a__						:	Optional[Any] =				self.num_labels
          a__						:	Optional[int] =				MegatronBertForSequenceClassification(lowercase)
          model.to(lowercase)
          model.eval()
          a__						:	int =				model(lowercase					,  attention_mask=lowercase					,  token_type_ids=lowercase					,  labels=lowercase)
          self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.num_labels))

   def 							__lowercase		(					self					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase)			->						Optional[int]:







          '''simple docstring'''

          a__						:	Any =				self.num_labels
          a__						:	Dict =				MegatronBertForTokenClassification(config=lowercase)
          model.to(lowercase)
          model.eval()
          a__						:	Dict =				model(lowercase					,  attention_mask=lowercase					,  token_type_ids=lowercase					,  labels=lowercase)
          self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.seq_length, self.num_labels))

   def 							__lowercase		(					self					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase					,  lowercase)			->						Dict:







          '''simple docstring'''

          a__						:	Any =				self.num_choices
          a__						:	Any =				MegatronBertForMultipleChoice(config=lowercase)
          model.to(lowercase)
          model.eval()
          a__						:	Any =				input_ids.unsqueeze(1).expand(-1					,  self.num_choices					,  -1).contiguous()
          a__						:	Any =				token_type_ids.unsqueeze(1).expand(-1					,  self.num_choices					,  -1).contiguous()
          a__						:	List[Any] =				input_mask.unsqueeze(1).expand(-1					,  self.num_choices					,  -1).contiguous()
          a__						:	Dict =				model(
              lowercase					,  attention_mask=lowercase					,  token_type_ids=lowercase					,  labels=lowercase					,  )
          self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.num_choices))

   def 							__lowercase		(					self)			->						Tuple:







          '''simple docstring'''

          a__						:	Dict =				self.prepare_config_and_inputs()
          (
              (
              a__
          ) ,							(
              a__
          ) ,							(
              a__
          ) ,							(
              a__
          ) ,							(
              a__
          ) ,							(
              a__
          ) ,							(
              a__
          ) ,							
          )						:	Optional[int] =				config_and_inputs
          a__						:	Optional[int] =				{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
          return config, inputs_dict





@require_torch
class  A__    (							__UpperCAmelCase ,  __UpperCAmelCase ,  unittest.TestCase	):




   """simple docstring"""



   __A   :				str         =  (
       (
           MegatronBertModel,
           MegatronBertForMaskedLM,
           MegatronBertForCausalLM,
           MegatronBertForMultipleChoice,
           MegatronBertForNextSentencePrediction,
           MegatronBertForPreTraining,
           MegatronBertForQuestionAnswering,
           MegatronBertForSequenceClassification,
           MegatronBertForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   __A   :				List[Any]         =  (
       {
           '''feature-extraction''': MegatronBertModel,
           '''fill-mask''': MegatronBertForMaskedLM,
           '''question-answering''': MegatronBertForQuestionAnswering,
           '''text-classification''': MegatronBertForSequenceClassification,
           '''text-generation''': MegatronBertForCausalLM,
           '''token-classification''': MegatronBertForTokenClassification,
           '''zero-shot''': MegatronBertForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   __A   :				List[Any]         =  True
   # test_resize_embeddings = False
   __A   :				Tuple         =  False

   def 							__lowercase		(					self					,  lowercase					,  lowercase					,  lowercase=False)			->						Optional[int]:







          '''simple docstring'''

          a__						:	List[Any] =				super()._prepare_for_class(lowercase					,  lowercase					,  return_labels=lowercase)

          if return_labels:
                 if model_class in get_values(lowercase):
                        a__						:	List[Any] =				torch.zeros(
                            (self.model_tester.batch_size, self.model_tester.seq_length)					,  dtype=torch.long					,  device=lowercase)
                        a__						:	Union[str, Any] =				torch.zeros(
                            self.model_tester.batch_size					,  dtype=torch.long					,  device=lowercase)
          return inputs_dict

   def 							__lowercase		(					self)			->						Optional[Any]:







          '''simple docstring'''

          a__						:	int =				MegatronBertModelTester(self)
          a__						:	List[Any] =				ConfigTester(self					,  config_class=lowercase					,  hidden_size=37)

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







          '''simple docstring'''

          self.config_tester.run_common_tests()

   def 							__lowercase		(					self)			->						int:







          '''simple docstring'''

          a__						:	List[str] =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_model(*lowercase)

   def 							__lowercase		(					self)			->						int:







          '''simple docstring'''

          a__						:	Any =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowercase)

   def 							__lowercase		(					self)			->						Optional[int]:







          '''simple docstring'''

          a__						:	Any =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowercase)

   def 							__lowercase		(					self)			->						List[str]:







          '''simple docstring'''

          a__						:	Dict =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowercase)

   def 							__lowercase		(					self)			->						Dict:







          '''simple docstring'''

          a__						:	str =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowercase)

   def 							__lowercase		(					self)			->						List[str]:







          '''simple docstring'''

          a__						:	str =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowercase)

   def 							__lowercase		(					self)			->						Optional[Any]:







          '''simple docstring'''

          a__						:	str =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowercase)

   def 							__lowercase		(					self)			->						Any:







          '''simple docstring'''

          a__						:	Optional[Any] =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowercase)



def A_	(		A__				)  ->		Any:
       return torch.tensor(
           A__		,      dtype=torch.long		,      device=A__		,      )


lowercase      :       Any									=  1e-4





@require_torch
@require_sentencepiece
@require_tokenizers
class  A__    (							unittest.TestCase	):




   """simple docstring"""



   @slow
   @unittest.skip('Model is not available.')
   def 							__lowercase		(					self)			->						List[Any]:







          '''simple docstring'''

          a__						:	List[str] =				'nvidia/megatron-bert-uncased-345m'
          if "MYDIR" in os.environ:
                 a__						:	Optional[int] =				os.path.join(os.environ['MYDIR']					,  lowercase)
          a__						:	str =				MegatronBertModel.from_pretrained(lowercase)
          model.to(lowercase)
          model.half()
          a__						:	List[str] =				_long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]])
          with torch.no_grad():
                 a__						:	int =				model(lowercase)[0]
          a__						:	Optional[Any] =				torch.Size((1, 9, 1024))
          self.assertEqual(output.shape					,  lowercase)

          a__						:	Tuple =				[-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
          for ii in range(3):
                 for jj in range(3):
                        a__						:	Optional[Any] =				output[0, ii, jj]
                        a__						:	Tuple =				expected[3 * ii + jj]
                        a__						:	Union[str, Any] =				'ii={} jj={} a={} b={}'.format(lowercase					,  lowercase					,  lowercase					,  lowercase)
                        self.assertTrue(math.isclose(lowercase					,  lowercase					,  rel_tol=lowercase					,  abs_tol=lowercase)					,  msg=lowercase)




import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

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 transformers import (
	    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
	    MODEL_MAPPING,
	    LevitForImageClassification,
	    LevitForImageClassificationWithTeacher,
	    LevitModel,
	)
	from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image

	from transformers import LevitImageProcessor







class        A_   (			__lowerCamelCase   ):




	'''simple docstring'''


	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    self.config_class(**self.inputs_dict       )
		self.parent.assertTrue(hasattr(snake_case		,    'hidden_sizes'       )       )
		self.parent.assertTrue(hasattr(snake_case		,    'num_attention_heads'       )       )







class        A_   :




	'''simple docstring'''


	def __init__(       self		,    snake_case		,    snake_case=13		,    snake_case=64		,    snake_case=3		,    snake_case=3		,    snake_case=2		,    snake_case=1		,    snake_case=16		,    snake_case=[128, 256, 384]		,    snake_case=[4, 6, 8]		,    snake_case=[2, 3, 4]		,    snake_case=[16, 16, 16]		,    snake_case=0		,    snake_case=[2, 2, 2]		,    snake_case=[2, 2, 2]		,    snake_case=0.02		,    snake_case=True		,    snake_case=True		,    snake_case=2		,    ):
		lowercase       =    parent
		lowercase       =    batch_size
		lowercase       =    image_size
		lowercase       =    num_channels
		lowercase       =    kernel_size
		lowercase       =    stride
		lowercase       =    padding
		lowercase       =    hidden_sizes
		lowercase       =    num_attention_heads
		lowercase       =    depths
		lowercase       =    key_dim
		lowercase       =    drop_path_rate
		lowercase       =    patch_size
		lowercase       =    attention_ratio
		lowercase       =    mlp_ratio
		lowercase       =    initializer_range
		lowercase       =    [
		    ['Subsample', key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
		    ['Subsample', key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
		]
		lowercase       =    is_training
		lowercase       =    use_labels
		lowercase       =    num_labels
		lowercase       =    initializer_range





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

		lowercase       =    None
		if self.use_labels:
			lowercase       =    ids_tensor([self.batch_size]		,    self.num_labels       )

		lowercase       =    self.get_config()
		return config, pixel_values, labels





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		return LevitConfig(
		    image_size=self.image_size		,    num_channels=self.num_channels		,    kernel_size=self.kernel_size		,    stride=self.stride		,    padding=self.padding		,    patch_size=self.patch_size		,    hidden_sizes=self.hidden_sizes		,    num_attention_heads=self.num_attention_heads		,    depths=self.depths		,    key_dim=self.key_dim		,    drop_path_rate=self.drop_path_rate		,    mlp_ratio=self.mlp_ratio		,    attention_ratio=self.attention_ratio		,    initializer_range=self.initializer_range		,    down_ops=self.down_ops		,    )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case       ):
		lowercase       =    LevitModel(config=snake_case       )
		model.to(snake_case       )
		model.eval()
		lowercase       =    model(snake_case       )
		lowercase       =    (self.image_size, self.image_size)
		lowercase  ,					lowercase       =    image_size[0], image_size[1]
		for _ in range(4       ):
			lowercase       =    floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1       )
			lowercase       =    floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1       )
		self.parent.assertEqual(
		    result.last_hidden_state.shape		,    (self.batch_size, ceil(height / 4       ) * ceil(width / 4       ), self.hidden_sizes[-1])		,    )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case       ):
		lowercase       =    self.num_labels
		lowercase       =    LevitForImageClassification(snake_case       )
		model.to(snake_case       )
		model.eval()
		lowercase       =    model(snake_case		,    labels=snake_case       )
		self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.num_labels)       )





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    self.prepare_config_and_inputs()
		lowercase  ,					lowercase  ,					lowercase       =    config_and_inputs
		lowercase       =    {'pixel_values': pixel_values}
		return config, inputs_dict







@require_torch
class        A_   (			__lowerCamelCase      ,	__lowerCamelCase      ,	unittest.TestCase   ):




	'''simple docstring'''


	_UpperCamelCase     :     Tuple							 =       (
	    (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
	    if is_torch_available()
	    else ()
	)
	_UpperCamelCase     :     Dict							 =       (
	    {
	        """feature-extraction""": LevitModel,
	        """image-classification""": (LevitForImageClassification, LevitForImageClassificationWithTeacher),
	    }
	    if is_torch_available()
	    else {}
	)

	_UpperCamelCase     :     Dict							 =       False
	_UpperCamelCase     :     List[str]							 =       False
	_UpperCamelCase     :     List[str]							 =       False
	_UpperCamelCase     :     str							 =       False
	_UpperCamelCase     :     List[str]							 =       False





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    LevitModelTester(self       )
		lowercase       =    ConfigTester(self		,    config_class=snake_case		,    has_text_modality=snake_case		,    hidden_size=37       )





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		self.create_and_test_config_common_properties()
		self.config_tester.create_and_test_config_to_json_string()
		self.config_tester.create_and_test_config_to_json_file()
		self.config_tester.create_and_test_config_from_and_save_pretrained()
		self.config_tester.create_and_test_config_with_num_labels()
		self.config_tester.check_config_can_be_init_without_params()
		self.config_tester.check_config_arguments_init()





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		return





	@unittest.skip(reason='Levit does not use inputs_embeds'       )
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		pass





	@unittest.skip(reason='Levit does not support input and output embeddings'       )
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		pass





	@unittest.skip(reason='Levit does not output attentions'       )
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		pass





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase  ,					lowercase       =    self.model_tester.prepare_config_and_inputs_for_common()

		for model_class in self.all_model_classes:
			lowercase       =    model_class(snake_case       )
			lowercase       =    inspect.signature(model.forward       )
			# signature.parameters is an OrderedDict => so arg_names order is deterministic
			lowercase       =    [*signature.parameters.keys()]

			lowercase       =    ['pixel_values']
			self.assertListEqual(arg_names[:1]		,    snake_case       )





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		def check_hidden_states_output(snake_case		,    snake_case		,    snake_case       ):
			lowercase       =    model_class(snake_case       )
			model.to(snake_case       )
			model.eval()

			with torch.no_grad():
				lowercase       =    model(**self._prepare_for_class(snake_case		,    snake_case       )       )

			lowercase       =    outputs.hidden_states

			lowercase       =    len(self.model_tester.depths       ) + 1
			self.assertEqual(len(snake_case       )		,    snake_case       )

			lowercase       =    (self.model_tester.image_size, self.model_tester.image_size)
			lowercase  ,					lowercase       =    image_size[0], image_size[1]
			for _ in range(4       ):
				lowercase       =    floor(
				    (
				        (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
				        / self.model_tester.stride
				    )
				    + 1       )
				lowercase       =    floor(
				    (
				        (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
				        / self.model_tester.stride
				    )
				    + 1       )
			# verify the first hidden states (first block)
			self.assertListEqual(
			    list(hidden_states[0].shape[-2:]       )		,    [
			        height * width,
			        self.model_tester.hidden_sizes[0],
			    ]		,    )

		lowercase  ,					lowercase       =    self.model_tester.prepare_config_and_inputs_for_common()

		for model_class in self.all_model_classes:
			lowercase       =    True
			check_hidden_states_output(snake_case		,    snake_case		,    snake_case       )

			# check that output_hidden_states also work using config
			del inputs_dict["output_hidden_states"]
			lowercase       =    True

			check_hidden_states_output(snake_case		,    snake_case		,    snake_case       )





	@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'       )
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		pass





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case=False       ):
		lowercase       =    super()._prepare_for_class(snake_case		,    snake_case		,    return_labels=snake_case       )

		if return_labels:
			if model_class.__name__ == "LevitForImageClassificationWithTeacher":
				del inputs_dict["labels"]

		return inputs_dict





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_model(*snake_case       )





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_image_classification(*snake_case       )





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		if not self.model_tester.is_training:
			return

		lowercase  ,					lowercase       =    self.model_tester.prepare_config_and_inputs_for_common()
		lowercase       =    True

		for model_class in self.all_model_classes:
			# LevitForImageClassificationWithTeacher supports inference-only
			if (
			    model_class in get_values(snake_case       )
			    or model_class.__name__ == "LevitForImageClassificationWithTeacher"
			):
				continue
			lowercase       =    model_class(snake_case       )
			model.to(snake_case       )
			model.train()
			lowercase       =    self._prepare_for_class(snake_case		,    snake_case		,    return_labels=snake_case       )
			lowercase       =    model(**snake_case       ).loss
			loss.backward()





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase  ,					lowercase       =    self.model_tester.prepare_config_and_inputs_for_common()
		if not self.model_tester.is_training:
			return

		lowercase       =    False
		lowercase       =    True

		for model_class in self.all_model_classes:
			if model_class in get_values(snake_case       ) or not model_class.supports_gradient_checkpointing:
				continue
			# LevitForImageClassificationWithTeacher supports inference-only
			if model_class.__name__ == "LevitForImageClassificationWithTeacher":
				continue
			lowercase       =    model_class(snake_case       )
			model.gradient_checkpointing_enable()
			model.to(snake_case       )
			model.train()
			lowercase       =    self._prepare_for_class(snake_case		,    snake_case		,    return_labels=snake_case       )
			lowercase       =    model(**snake_case       ).loss
			loss.backward()





	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase  ,					lowercase       =    self.model_tester.prepare_config_and_inputs_for_common()

		lowercase       =    [
		    {'title': 'multi_label_classification', 'num_labels': 2, 'dtype': torch.float},
		    {'title': 'single_label_classification', 'num_labels': 1, 'dtype': torch.long},
		    {'title': 'regression', 'num_labels': 1, 'dtype': torch.float},
		]

		for model_class in self.all_model_classes:
			if (
			    model_class
			    not in [
			        *get_values(snake_case       ),
			    ]
			    or model_class.__name__ == "LevitForImageClassificationWithTeacher"
			):
				continue

			for problem_type in problem_types:
				with self.subTest(msg=F'''Testing {model_class} with {problem_type['title']}'''       ):
					lowercase       =    problem_type['title']
					lowercase       =    problem_type['num_labels']

					lowercase       =    model_class(snake_case       )
					model.to(snake_case       )
					model.train()

					lowercase       =    self._prepare_for_class(snake_case		,    snake_case		,    return_labels=snake_case       )

					if problem_type["num_labels"] > 1:
						lowercase       =    inputs['labels'].unsqueeze(1       ).repeat(1		,    problem_type['num_labels']       )

					lowercase       =    inputs['labels'].to(problem_type['dtype']       )

					# This tests that we do not trigger the warning form PyTorch "Using a target size that is different
					# to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
					# they have the same size." which is a symptom something in wrong for the regression problem.
					# See https://github.com/huggingface/transformers/issues/11780
					with warnings.catch_warnings(record=snake_case       ) as warning_list:
						lowercase       =    model(**snake_case       ).loss
					for w in warning_list:
						if "Using a target size that is different to the input size" in str(w.message       ):
							raise ValueError(
							    F'''Something is going wrong in the regression problem: intercepted {w.message}'''       )

					loss.backward()





	@slow
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
			lowercase       =    LevitModel.from_pretrained(snake_case       )
			self.assertIsNotNone(snake_case       )






def 							UpperCAmelCase_		(							):
	lowercase       =    Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'							)
	return image







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




	'''simple docstring'''


	@cached_property
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]       )





	@slow
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		lowercase       =    LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]       ).to(
		    snake_case       )

		lowercase       =    self.default_image_processor
		lowercase       =    prepare_img()
		lowercase       =    image_processor(images=snake_case		,    return_tensors='pt'       ).to(snake_case       )

		# forward pass
		with torch.no_grad():
			lowercase       =    model(**snake_case       )

		# verify the logits
		lowercase       =    torch.Size((1, 1000)       )
		self.assertEqual(outputs.logits.shape		,    snake_case       )

		lowercase       =    torch.tensor([1.0_448, -0.3_745, -1.8_317]       ).to(snake_case       )

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

