


'''simple docstring'''



import inspect
import unittest
import warnings

from transformers import DeiTConfig
from transformers.models.auto import get_values
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 (
		    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
		    MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
		    MODEL_MAPPING,
		    DeiTForImageClassification,
		    DeiTForImageClassificationWithTeacher,
		    DeiTForMaskedImageModeling,
		    DeiTModel,
		)
		from transformers.models.deit.modeling_deit import DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import DeiTImageProcessor


class 		_snake_case					:




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase=13							,							_lowerCamelCase=30							,							_lowerCamelCase=2							,							_lowerCamelCase=3							,							_lowerCamelCase=True							,							_lowerCamelCase=True							,							_lowerCamelCase=32							,							_lowerCamelCase=5							,							_lowerCamelCase=4							,							_lowerCamelCase=37							,							_lowerCamelCase="gelu"							,							_lowerCamelCase=0.1							,							_lowerCamelCase=0.1							,							_lowerCamelCase=10							,							_lowerCamelCase=0.02							,							_lowerCamelCase=3							,							_lowerCamelCase=None							,							_lowerCamelCase=2							,							):
								UpperCAmelCase__  :					Any        =     parent
								UpperCAmelCase__  :					Optional[Any]        =     batch_size
								UpperCAmelCase__  :					List[str]        =     image_size
								UpperCAmelCase__  :					int        =     patch_size
								UpperCAmelCase__  :					Union[str, Any]        =     num_channels
								UpperCAmelCase__  :					Tuple        =     is_training
								UpperCAmelCase__  :					str        =     use_labels
								UpperCAmelCase__  :					List[str]        =     hidden_size
								UpperCAmelCase__  :					Union[str, Any]        =     num_hidden_layers
								UpperCAmelCase__  :					Optional[Any]        =     num_attention_heads
								UpperCAmelCase__  :					Any        =     intermediate_size
								UpperCAmelCase__  :					Union[str, Any]        =     hidden_act
								UpperCAmelCase__  :					Tuple        =     hidden_dropout_prob
								UpperCAmelCase__  :					List[Any]        =     attention_probs_dropout_prob
								UpperCAmelCase__  :					str        =     type_sequence_label_size
								UpperCAmelCase__  :					List[Any]        =     initializer_range
								UpperCAmelCase__  :					List[Any]        =     scope
								UpperCAmelCase__  :					Any        =     encoder_stride

								# in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
								UpperCAmelCase__  :					int        =     (image_size // patch_size) ** 2
								UpperCAmelCase__  :					Optional[Any]        =     num_patches + 2




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Optional[int]        =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

								UpperCAmelCase__  :					Union[str, Any]        =     self.get_config()

								return config, pixel_values, labels




					def 				snake_case__						(      self):
								return DeiTConfig(
								    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=lowerCamelCase_							,							initializer_range=self.initializer_range							,							encoder_stride=self.encoder_stride							,							)




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase):
								UpperCAmelCase__  :					Any        =     DeiTModel(config=lowerCamelCase_)
								model.to(lowerCamelCase_)
								model.eval()
								UpperCAmelCase__  :					Optional[Any]        =     model(lowerCamelCase_)
								self.parent.assertEqual(result.last_hidden_state.shape							,							(self.batch_size, self.seq_length, self.hidden_size))




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase):
								UpperCAmelCase__  :					str        =     DeiTForMaskedImageModeling(config=lowerCamelCase_)
								model.to(lowerCamelCase_)
								model.eval()
								UpperCAmelCase__  :					Union[str, Any]        =     model(lowerCamelCase_)
								self.parent.assertEqual(
								    result.reconstruction.shape							,							(self.batch_size, self.num_channels, self.image_size, self.image_size))

								# test greyscale images
								UpperCAmelCase__  :					str        =     1
								UpperCAmelCase__  :					List[Any]        =     DeiTForMaskedImageModeling(lowerCamelCase_)
								model.to(lowerCamelCase_)
								model.eval()

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




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase):
								UpperCAmelCase__  :					Optional[Any]        =     self.type_sequence_label_size
								UpperCAmelCase__  :					Union[str, Any]        =     DeiTForImageClassification(lowerCamelCase_)
								model.to(lowerCamelCase_)
								model.eval()
								UpperCAmelCase__  :					int        =     model(lowerCamelCase_							,							labels=lowerCamelCase_)
								self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.type_sequence_label_size))

								# test greyscale images
								UpperCAmelCase__  :					List[Any]        =     1
								UpperCAmelCase__  :					Tuple        =     DeiTForImageClassification(lowerCamelCase_)
								model.to(lowerCamelCase_)
								model.eval()

								UpperCAmelCase__  :					Tuple        =     floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
								UpperCAmelCase__  :					Optional[Any]        =     model(lowerCamelCase_							,							labels=lowerCamelCase_)
								self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.type_sequence_label_size))




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Optional[int]        =     self.prepare_config_and_inputs()
								(
								    UpperCAmelCase__					
								)  :					List[str]        =     config_and_inputs
								UpperCAmelCase__  :					Any        =     {'''pixel_values''': pixel_values}
								return config, inputs_dict


@require_torch
class 		_snake_case					(     _UpperCAmelCase       ,   _UpperCAmelCase       ,   unittest.TestCase	):
					lowerCAmelCase      :Union[str, Any]   =					(
					    (
					        DeiTModel,
					        DeiTForImageClassification,
					        DeiTForImageClassificationWithTeacher,
					        DeiTForMaskedImageModeling,
					    )
					    if is_torch_available()
					    else ()
					)
					lowerCAmelCase      :List[Any]   =					(
					    {
					        '''feature-extraction''': DeiTModel,
					        '''image-classification''': (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
					    }
					    if is_torch_available()
					    else {}
					)

					lowerCAmelCase      :Union[str, Any]   =					False
					lowerCAmelCase      :Tuple   =					False
					lowerCAmelCase      :Optional[Any]   =					False




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Dict        =     DeiTModelTester(self)
								UpperCAmelCase__  :					Union[str, Any]        =     ConfigTester(self							,							config_class=lowerCamelCase_							,							has_text_modality=lowerCamelCase_							,							hidden_size=37)




					def 				snake_case__						(      self):
								self.config_tester.run_common_tests()




					@unittest.skip(reason="""DeiT does not use inputs_embeds""")
					def 				snake_case__						(      self):
								pass




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					str        =     self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											UpperCAmelCase__  :					Dict        =     model_class(lowerCamelCase_)
											self.assertIsInstance(model.get_input_embeddings()							,							(nn.Module))
											UpperCAmelCase__  :					Dict        =     model.get_output_embeddings()
											self.assertTrue(x is None or isinstance(lowerCamelCase_							,							nn.Linear))




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					List[Any]        =     self.model_tester.prepare_config_and_inputs_for_common()

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

											UpperCAmelCase__  :					Tuple        =     ['''pixel_values''']
											self.assertListEqual(arg_names[:1]							,							lowerCamelCase_)




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					str        =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCamelCase_)




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					List[Any]        =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_image_modeling(*lowerCamelCase_)




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Optional[Any]        =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*lowerCamelCase_)




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase=False):
								UpperCAmelCase__  :					Optional[Any]        =     super()._prepare_for_class(lowerCamelCase_							,							lowerCamelCase_							,							return_labels=lowerCamelCase_)

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

								return inputs_dict




					def 				snake_case__						(      self):




								if not self.model_tester.is_training:
											return

								UpperCAmelCase__  :					Any        =     self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase__  :					str        =     True

								for model_class in self.all_model_classes:
											# DeiTForImageClassificationWithTeacher supports inference-only
											if (
											    model_class in get_values(lowerCamelCase_)
											    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
											):
														continue
											UpperCAmelCase__  :					int        =     model_class(lowerCamelCase_)
											model.to(lowerCamelCase_)
											model.train()
											UpperCAmelCase__  :					Tuple        =     self._prepare_for_class(lowerCamelCase_							,							lowerCamelCase_							,							return_labels=lowerCamelCase_)
											UpperCAmelCase__  :					Any        =     model(**lowerCamelCase_).loss
											loss.backward()




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					List[str]        =     self.model_tester.prepare_config_and_inputs_for_common()
								if not self.model_tester.is_training:
											return

								UpperCAmelCase__  :					Union[str, Any]        =     False
								UpperCAmelCase__  :					Optional[Any]        =     True

								for model_class in self.all_model_classes:
											if model_class in get_values(lowerCamelCase_) or not model_class.supports_gradient_checkpointing:
														continue
											# DeiTForImageClassificationWithTeacher supports inference-only
											if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
														continue
											UpperCAmelCase__  :					Dict        =     model_class(lowerCamelCase_)
											model.gradient_checkpointing_enable()
											model.to(lowerCamelCase_)
											model.train()
											UpperCAmelCase__  :					Optional[int]        =     self._prepare_for_class(lowerCamelCase_							,							lowerCamelCase_							,							return_labels=lowerCamelCase_)
											UpperCAmelCase__  :					Tuple        =     model(**lowerCamelCase_).loss
											loss.backward()




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					int        =     self.model_tester.prepare_config_and_inputs_for_common()

								UpperCAmelCase__  :					Tuple        =     [
								    {'''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(lowerCamelCase_),
											        *get_values(lowerCamelCase_),
											    ]
											    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
											):
														continue

											for problem_type in problem_types:
														with self.subTest(msg=f'''Testing {model_class} with {problem_type['title']}'''):
																	UpperCAmelCase__  :					str        =     problem_type['''title''']
																	UpperCAmelCase__  :					Dict        =     problem_type['''num_labels''']

																	UpperCAmelCase__  :					List[Any]        =     model_class(lowerCamelCase_)
																	model.to(lowerCamelCase_)
																	model.train()

																	UpperCAmelCase__  :					str        =     self._prepare_for_class(lowerCamelCase_							,							lowerCamelCase_							,							return_labels=lowerCamelCase_)

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

																	UpperCAmelCase__  :					Union[str, Any]        =     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=lowerCamelCase_) as warning_list:
																				UpperCAmelCase__  :					List[str]        =     model(**lowerCamelCase_).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 				snake_case__						(      self):




								for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											UpperCAmelCase__  :					List[str]        =     DeiTModel.from_pretrained(lowerCamelCase_)
											self.assertIsNotNone(lowerCamelCase_)







def     _UpperCamelCase      (		):
			UpperCAmelCase__  :					str        =     Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""					)
			return image


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




					@cached_property
					def 				snake_case__						(      self):
								return (
								    DeiTImageProcessor.from_pretrained("""facebook/deit-base-distilled-patch16-224""")
								    if is_vision_available()
								    else None
								)




					@slow
					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Tuple        =     DeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224""").to(
								    lowerCamelCase_)

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

								# forward pass
								with torch.no_grad():
											UpperCAmelCase__  :					Tuple        =     model(**lowerCamelCase_)

								# verify the logits
								UpperCAmelCase__  :					str        =     torch.Size((1, 1000))
								self.assertEqual(outputs.logits.shape							,							lowerCamelCase_)

								UpperCAmelCase__  :					Tuple        =     torch.tensor([-1.0266, 0.1912, -1.2861]).to(lowerCamelCase_)

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




					@slow
					@require_accelerate
					@require_torch_gpu
					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Optional[Any]        =     DeiTModel.from_pretrained(
								    """facebook/deit-base-distilled-patch16-224"""							,							torch_dtype=torch.floataa							,							device_map="""auto""")
								UpperCAmelCase__  :					Optional[int]        =     self.default_image_processor

								UpperCAmelCase__  :					List[str]        =     prepare_img()
								UpperCAmelCase__  :					List[Any]        =     image_processor(images=lowerCamelCase_							,							return_tensors="""pt""")
								UpperCAmelCase__  :					Tuple        =     inputs.pixel_values.to(lowerCamelCase_)

								# forward pass to make sure inference works in fp16
								with torch.no_grad():
											UpperCAmelCase__  :					List[Any]        =     model(lowerCamelCase_)



'''simple docstring'''



from typing import Callable, Dict, Optional, Tuple

import torch
from torch import nn
from torch.distributions import (
    AffineTransform,
    Distribution,
    Independent,
    NegativeBinomial,
    Normal,
    StudentT,
    TransformedDistribution,
)


class 		_snake_case					(     a__	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase=None							,							_lowerCamelCase=None							,							_lowerCamelCase=0):
								UpperCAmelCase__  :					Dict        =     1.0 if scale is None else scale
								UpperCAmelCase__  :					Dict        =     0.0 if loc is None else loc

								super().__init__(_lowerCamelCase							,							[AffineTransform(loc=self.loc							,							scale=self.scale							,							event_dim=_lowerCamelCase)])




					@property
					def 				snake_case__						(      self):
								return self.base_dist.mean * self.scale + self.loc




					@property
					def 				snake_case__						(      self):
								return self.base_dist.variance * self.scale**2




					@property
					def 				snake_case__						(      self):
								return self.variance.sqrt()


class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase							,							**_lowerCamelCase):
								super().__init__(**_lowerCamelCase)
								UpperCAmelCase__  :					int        =     args_dim
								UpperCAmelCase__  :					Optional[int]        =     nn.ModuleList([nn.Linear(_lowerCamelCase							,							_lowerCamelCase) for dim in args_dim.values()])
								UpperCAmelCase__  :					Optional[Any]        =     domain_map




					def 				snake_case__						(      self							,							_lowerCamelCase):
								UpperCAmelCase__  :					Optional[int]        =     [proj(_lowerCamelCase) for proj in self.proj]

								return self.domain_map(*_lowerCamelCase)


class 		_snake_case					(     nn.Module	):




					def __init__(      self							,							_lowerCamelCase):
								super().__init__()
								UpperCAmelCase__  :					Optional[int]        =     function




					def 				snake_case__						(      self							,							_lowerCamelCase							,							*_lowerCamelCase):
								return self.function(_lowerCamelCase							,							*_lowerCamelCase)


class 		_snake_case					:
					lowerCAmelCase      :type
					lowerCAmelCase      :int
					lowerCAmelCase      :Dict[str, int]




					def __init__(      self							,							_lowerCamelCase = 1):
								UpperCAmelCase__  :					Optional[Any]        =     dim
								UpperCAmelCase__  :					int        =     {k: dim * self.args_dim[k] for k in self.args_dim}




					def 				snake_case__						(      self							,							_lowerCamelCase):
								if self.dim == 1:
											return self.distribution_class(*_lowerCamelCase)
								else:
											return Independent(self.distribution_class(*_lowerCamelCase)							,							1)




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase = None							,							_lowerCamelCase = None							,							):
								UpperCAmelCase__  :					Dict        =     self._base_distribution(_lowerCamelCase)
								if loc is None and scale is None:
											return distr
								else:
											return AffineTransformed(_lowerCamelCase							,							loc=_lowerCamelCase							,							scale=_lowerCamelCase							,							event_dim=self.event_dim)




					@property
					def 				snake_case__						(      self):
								return () if self.dim == 1 else (self.dim,)




					@property
					def 				snake_case__						(      self):
								return len(self.event_shape)




					@property
					def 				snake_case__						(      self):
								return 0.0




					def 				snake_case__						(      self							,							_lowerCamelCase):
								return ParameterProjection(
								    in_features=_lowerCamelCase							,							args_dim=self.args_dim							,							domain_map=LambdaLayer(self.domain_map)							,							)




					def 				snake_case__						(      self							,							*_lowerCamelCase):
								raise NotImplementedError()




					@staticmethod
					def 				snake_case__						(      _lowerCamelCase):
								return (x + torch.sqrt(torch.square(_lowerCamelCase) + 4.0)) / 2.0


class 		_snake_case					(     a__	):
					lowerCAmelCase      :Dict[str, int]   =					{"df": 1, "loc": 1, "scale": 1}
					lowerCAmelCase      :type   =					StudentT




					@classmethod
					def 				snake_case__						(      cls							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase):
								UpperCAmelCase__  :					Optional[Any]        =     cls.squareplus(_lowerCamelCase).clamp_min(torch.finfo(scale.dtype).eps)
								UpperCAmelCase__  :					Any        =     2.0 + cls.squareplus(_lowerCamelCase)
								return df.squeeze(-1), loc.squeeze(-1), scale.squeeze(-1)


class 		_snake_case					(     a__	):
					lowerCAmelCase      :Dict[str, int]   =					{"loc": 1, "scale": 1}
					lowerCAmelCase      :type   =					Normal




					@classmethod
					def 				snake_case__						(      cls							,							_lowerCamelCase							,							_lowerCamelCase):
								UpperCAmelCase__  :					str        =     cls.squareplus(_lowerCamelCase).clamp_min(torch.finfo(scale.dtype).eps)
								return loc.squeeze(-1), scale.squeeze(-1)


class 		_snake_case					(     a__	):
					lowerCAmelCase      :Dict[str, int]   =					{"total_count": 1, "logits": 1}
					lowerCAmelCase      :type   =					NegativeBinomial




					@classmethod
					def 				snake_case__						(      cls							,							_lowerCamelCase							,							_lowerCamelCase):
								UpperCAmelCase__  :					str        =     cls.squareplus(_lowerCamelCase)
								return total_count.squeeze(-1), logits.squeeze(-1)




					def 				snake_case__						(      self							,							_lowerCamelCase):
								UpperCAmelCase__   ,       UpperCAmelCase__  :					Tuple        =     distr_args
								if self.dim == 1:
											return self.distribution_class(total_count=_lowerCamelCase							,							logits=_lowerCamelCase)
								else:
											return Independent(self.distribution_class(total_count=_lowerCamelCase							,							logits=_lowerCamelCase)							,							1)




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase = None							,							_lowerCamelCase = None):
								UpperCAmelCase__   ,       UpperCAmelCase__  :					Any        =     distr_args

								if scale is not None:
											# See scaling property of Gamma.
											logits += scale.log()

								return self._base_distribution((total_count, logits))