


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





     '''simple docstring'''



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

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

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

       snake_case__						= None
       if self.use_labels:
         snake_case__						= ids_tensor([self.batch_size]     ,		self.type_sequence_label_size   )

       snake_case__						= self.get_config()

       return config, pixel_values, labels

     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]   ):
       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=_a     ,		initializer_range=self.initializer_range     ,		encoder_stride=self.encoder_stride     ,		)

     def 	SCREAMING_SNAKE_CASE__						(					self:Dict     ,		_a:Tuple     ,		_a:Any     ,		_a:Optional[int]   ):
       snake_case__						= ViTModel(config=_a   )
       model.to(_a   )
       model.eval()
       snake_case__						= model(_a   )
       self.parent.assertEqual(result.last_hidden_state.shape     ,		(self.batch_size, self.seq_length, self.hidden_size)   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[Any]     ,		_a:Optional[Any]     ,		_a:Dict     ,		_a:Union[str, Any]   ):
       snake_case__						= ViTForMaskedImageModeling(config=_a   )
       model.to(_a   )
       model.eval()
       snake_case__						= model(_a   )
       self.parent.assertEqual(
           result.reconstruction.shape     ,		(self.batch_size, self.num_channels, self.image_size, self.image_size)   )

       # test greyscale images
       snake_case__						= 1
       snake_case__						= ViTForMaskedImageModeling(_a   )
       model.to(_a   )
       model.eval()

       snake_case__						= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
       snake_case__						= model(_a   )
       self.parent.assertEqual(result.reconstruction.shape     ,		(self.batch_size, 1, self.image_size, self.image_size)   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Dict     ,		_a:Tuple     ,		_a:str     ,		_a:Optional[int]   ):
       snake_case__						= self.type_sequence_label_size
       snake_case__						= ViTForImageClassification(_a   )
       model.to(_a   )
       model.eval()
       snake_case__						= model(_a     ,		labels=_a   )
       self.parent.assertEqual(result.logits.shape     ,		(self.batch_size, self.type_sequence_label_size)   )

       # test greyscale images
       snake_case__						= 1
       snake_case__						= ViTForImageClassification(_a   )
       model.to(_a   )
       model.eval()

       snake_case__						= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
       snake_case__						= model(_a   )
       self.parent.assertEqual(result.logits.shape     ,		(self.batch_size, self.type_sequence_label_size)   )







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




@require_torch
class 		__magic_name__						(snake_case_   ,snake_case_   ,unittest.TestCase				):





     '''simple docstring'''



     __lowercase				:							Tuple	   =       (
         (
             ViTModel,
             ViTForImageClassification,
             ViTForMaskedImageModeling,
         )
         if is_torch_available()
         else ()
     )
     __lowercase				:							Dict	   =       (
         {'feature-extraction': ViTModel, 'image-classification': ViTForImageClassification}
         if is_torch_available()
         else {}
     )
     __lowercase				:							Any	   =       True

     __lowercase				:							Union[str, Any]	   =       False
     __lowercase				:							int	   =       False
     __lowercase				:							List[str]	   =       False

     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       snake_case__						= ViTModelTester(self   )
       snake_case__						= ConfigTester(self     ,		config_class=_a     ,		has_text_modality=_a     ,		hidden_size=37   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Dict   ):
       self.config_tester.run_common_tests()

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

     def 	SCREAMING_SNAKE_CASE__						(					self:Any   ):
       snake_case__				,				snake_case__						= self.model_tester.prepare_config_and_inputs_for_common()

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

     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[Any]   ):
       snake_case__				,				snake_case__						= self.model_tester.prepare_config_and_inputs_for_common()

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

         snake_case__						= ['''pixel_values''']
         self.assertListEqual(arg_names[:1]     ,		_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       snake_case__						= self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Union[str, Any]   ):
       snake_case__						= self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_image_modeling(*_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple   ):
       snake_case__						= self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*_a   )







     @slow
     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple   ):
       for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case__						= ViTModel.from_pretrained(_a   )
         self.assertIsNotNone(_a   )







def SCREAMING_SNAKE_CASE	(			)    ->				Union[str, Any]:
  snake_case__						= Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
  return image


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





     '''simple docstring'''



     @cached_property
     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]   ):
       return ViTImageProcessor.from_pretrained('''google/vit-base-patch16-224'''   ) if is_vision_available() else None

     @slow
     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       snake_case__						= ViTForImageClassification.from_pretrained('''google/vit-base-patch16-224'''   ).to(_a   )

       snake_case__						= self.default_image_processor
       snake_case__						= prepare_img()
       snake_case__						= image_processor(images=_a     ,		return_tensors='''pt'''   ).to(_a   )

       # forward pass
       with torch.no_grad():
         snake_case__						= model(**_a   )

       # verify the logits
       snake_case__						= torch.Size((1, 10_00)   )
       self.assertEqual(outputs.logits.shape     ,		_a   )

       snake_case__						= torch.tensor([-0.2744, 0.8215, -0.0836]   ).to(_a   )

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

     @slow
     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       # ViT models have an `interpolate_pos_encoding` argument in their forward method,
       # allowing to interpolate the pre-trained position embeddings in order to use
       # the model on higher resolutions. The DINO model by Facebook AI leverages this
       # to visualize self-attention on higher resolution images.
       snake_case__						= ViTModel.from_pretrained('''facebook/dino-vits8'''   ).to(_a   )

       snake_case__						= ViTImageProcessor.from_pretrained('''facebook/dino-vits8'''     ,		size=4_80   )
       snake_case__						= prepare_img()
       snake_case__						= image_processor(images=_a     ,		return_tensors='''pt'''   )
       snake_case__						= inputs.pixel_values.to(_a   )

       # forward pass
       with torch.no_grad():
         snake_case__						= model(_a     ,		interpolate_pos_encoding=_a   )

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

       snake_case__						= torch.tensor(
           [[4.2340, 4.3906, -6.6692], [4.5463, 1.8928, -6.7257], [4.4429, 0.8496, -5.8585]]   ).to(_a   )

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







     @slow
     @require_accelerate
     @require_torch_gpu
     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[Any]   ):
       snake_case__						= ViTModel.from_pretrained('''facebook/dino-vits8'''     ,		torch_dtype=torch.floataa     ,		device_map='''auto'''   )
       snake_case__						= self.default_image_processor

       snake_case__						= prepare_img()
       snake_case__						= image_processor(images=_a     ,		return_tensors='''pt'''   )
       snake_case__						= inputs.pixel_values.to(_a   )

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



def  __snake_case							(	__magic_name__		,    __magic_name__ = False						):







				'''simple docstring'''


				if n == 2:
								return True
				if not n % 2 or n < 2:
								return False
				if n > 5 and n % 10 not in (1, 3, 7, 9):  # can quickly check last digit
								return False
				if n > 3317044064679887385961981 and not allow_probable:
								raise ValueError(
								    "Warning: upper bound of deterministic test is exceeded. "
								    "Pass allow_probable=True to allow probabilistic test. "
								    "A return value of True indicates a probable prime."						)
				# array bounds provided by analysis
				lowercase         =      [
				    2047,
				    1373653,
				    25326001,
				    3215031751,
				    2152302898747,
				    3474749660383,
				    341550071728321,
				    1,
				    3825123056546413051,
				    1,
				    1,
				    318665857834031151167461,
				    3317044064679887385961981,
				]

				lowercase         =      [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41]
				for idx, _p in enumerate(__magic_name__		,    1						):
								if n < _p:
												# then we have our last prime to check
												lowercase         =      primes[:idx]
												break
				lowercase     ,	lowercase         =      n - 1, 0
				# break up n -1 into a power of 2 (s) and
				# remaining odd component
				# essentially, solve for d * 2 ** s == n - 1
				while d % 2 == 0:
								d //= 2
								s += 1
				for prime in plist:
								lowercase         =      False
								for r in range(__magic_name__						):
												lowercase         =      pow(__magic_name__		,    d * 2**r		,    __magic_name__						)
												# see article for analysis explanation for m
												if (r == 0 and m == 1) or ((m + 1) % n == 0):
																lowercase         =      True
																# this loop will not determine compositeness
																break
								if pr:
												continue
								# if pr is False, then the above loop never evaluated to true,
								# and the n MUST be composite
								return False
				return True

def  __snake_case							(	):







				'''simple docstring'''


				assert not miller_rabin(561						)
				assert miller_rabin(563						)
				# 2047

				assert not miller_rabin(838201						)
				assert miller_rabin(838207						)
				# 1_373_653

				assert not miller_rabin(17316001						)
				assert miller_rabin(17316017						)
				# 25_326_001

				assert not miller_rabin(3078386641						)
				assert miller_rabin(3078386653						)
				# 3_215_031_751

				assert not miller_rabin(1713045574801						)
				assert miller_rabin(1713045574819						)
				# 2_152_302_898_747

				assert not miller_rabin(2779799728307						)
				assert miller_rabin(2779799728327						)
				# 3_474_749_660_383

				assert not miller_rabin(113850023909441						)
				assert miller_rabin(113850023909527						)
				# 341_550_071_728_321

				assert not miller_rabin(1275041018848804351						)
				assert miller_rabin(1275041018848804391						)
				# 3_825_123_056_546_413_051

				assert not miller_rabin(79666464458507787791867						)
				assert miller_rabin(79666464458507787791951						)
				# 318_665_857_834_031_151_167_461

				assert not miller_rabin(552840677446647897660333						)
				assert miller_rabin(552840677446647897660359						)
				# 3_317_044_064_679_887_385_961_981
				# upper limit for probabilistic test


if __name__ == "__main__":
						test_miller_rabin()
