




'''simple docstring'''


import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
				import jax

				from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel

class     A					(	unittest.TestCase					):







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


														"""simple docstring"""


														_a			  =       parent
														_a			  =       batch_size
														_a			  =       image_size
														_a			  =       patch_size
														_a			  =       num_channels
														_a			  =       is_training
														_a			  =       use_labels
														_a			  =       hidden_size
														_a			  =       num_hidden_layers
														_a			  =       num_attention_heads
														_a			  =       intermediate_size
														_a			  =       hidden_act
														_a			  =       hidden_dropout_prob
														_a			  =       attention_probs_dropout_prob
														_a			  =       type_sequence_label_size
														_a			  =       initializer_range

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







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


														"""simple docstring"""


														_a			  =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

														_a			  =       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=lowerCAmelCase_							,    initializer_range=self.initializer_range							,    )

														return config, pixel_values







							def      __lowerCAmelCase    (						self :						Optional[Any]							,    lowerCAmelCase_ :						Union[str, Any]							,    lowerCAmelCase_ :						List[str]      )    ->							int:


														"""simple docstring"""


														_a			  =       FlaxViTModel(config=lowerCAmelCase_      )
														_a			  =       model(lowerCAmelCase_      )
														# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
														_a			  =       (self.image_size, self.image_size)
														_a			  =       (self.patch_size, self.patch_size)
														_a			  =       (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
														self.parent.assertEqual(result.last_hidden_state.shape							,    (self.batch_size, num_patches + 1, self.hidden_size)      )







							def      __lowerCAmelCase    (						self :						int							,    lowerCAmelCase_ :						Any							,    lowerCAmelCase_ :						List[str]      )    ->							Union[str, Any]:


														"""simple docstring"""


														_a			  =       self.type_sequence_label_size
														_a			  =       FlaxViTForImageClassification(config=lowerCAmelCase_      )
														_a			  =       model(lowerCAmelCase_      )
														self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.type_sequence_label_size)      )

														# test greyscale images
														_a			  =       1
														_a			  =       FlaxViTForImageClassification(lowerCAmelCase_      )

														_a			  =       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
														_a			  =       model(lowerCAmelCase_      )







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


														"""simple docstring"""


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



@require_flax
class     A					(	_a  ,unittest.TestCase					):
							lowercase_           =       (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()







							def      __lowerCAmelCase    (						self :						Any      )    ->							None:


														"""simple docstring"""


														_a			  =       FlaxViTModelTester(self      )
														_a			  =       ConfigTester(self							,    config_class=lowerCAmelCase_							,    has_text_modality=lowerCAmelCase_							,    hidden_size=37      )







							def      __lowerCAmelCase    (						self :						Tuple      )    ->							Any:


														"""simple docstring"""


														self.config_tester.run_common_tests()







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


														"""simple docstring"""


														_a			  =       self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*lowerCAmelCase_      )







							def      __lowerCAmelCase    (						self :						str      )    ->							List[Any]:


														"""simple docstring"""


														_a			  =       self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase_      )







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


														"""simple docstring"""


														_a  ,							_a			  =       self.model_tester.prepare_config_and_inputs_for_common()

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

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







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


														"""simple docstring"""


														_a  ,							_a			  =       self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					with self.subTest(model_class.__name__      ):
																												_a			  =       self._prepare_for_class(lowerCAmelCase_							,    lowerCAmelCase_      )
																												_a			  =       model_class(lowerCAmelCase_      )

																												@jax.jit
																												def model_jitted(lowerCAmelCase_ :						Union[str, Any]							,    **lowerCAmelCase_ :						Optional[Any]      ):
																																			return model(pixel_values=lowerCAmelCase_							,    **lowerCAmelCase_      )

																												with self.subTest('''JIT Enabled'''      ):
																																			_a			  =       model_jitted(**lowerCAmelCase_      ).to_tuple()

																												with self.subTest('''JIT Disabled'''      ):
																																			with jax.disable_jit():
																																										_a			  =       model_jitted(**lowerCAmelCase_      ).to_tuple()

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







							@slow
							def      __lowerCAmelCase    (						self :						str      )    ->							Any:


														"""simple docstring"""


														for model_class_name in self.all_model_classes:
																					_a			  =       model_class_name.from_pretrained('''google/vit-base-patch16-224'''      )
																					_a			  =       model(np.ones((1, 3, 2_24, 2_24)      )      )
																					self.assertIsNotNone(lowerCAmelCase_      )







'''simple docstring'''


import math
import unittest
def 	snake_case_					(UpperCamelCase :						int   ):





							'''simple docstring'''


							assert isinstance(UpperCamelCase				,		UpperCamelCase   ) and (
							    number >= 0
							), "'number' must been an int and positive"

							if 1 < number < 4:
														# 2 and 3 are primes
														return True
							elif number < 2 or number % 2 == 0 or number % 3 == 0:
														# Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
														return False

							# All primes number are in format of 6k +/- 1
							for i in range(5				,		int(math.sqrt(UpperCamelCase   ) + 1   )				,		6   ):
														if number % i == 0 or number % (i + 2) == 0:
																					return False
							return True

class     A					(	unittest.TestCase					):







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


														"""simple docstring"""


														self.assertTrue(is_prime(2      )      )
														self.assertTrue(is_prime(3      )      )
														self.assertTrue(is_prime(5      )      )
														self.assertTrue(is_prime(7      )      )
														self.assertTrue(is_prime(11      )      )
														self.assertTrue(is_prime(13      )      )
														self.assertTrue(is_prime(17      )      )
														self.assertTrue(is_prime(19      )      )
														self.assertTrue(is_prime(23      )      )
														self.assertTrue(is_prime(29      )      )







							def      __lowerCAmelCase    (						self :						Any      )    ->							Dict:


														"""simple docstring"""


														with self.assertRaises(lowerCAmelCase_      ):
																					is_prime(-19      )
														self.assertFalse(
														    is_prime(0      )							,    '''Zero doesn\'t have any positive factors, primes must have exactly two.'''							,    )
														self.assertFalse(
														    is_prime(1      )							,    '''One only has 1 positive factor, primes must have exactly two.'''							,    )
														self.assertFalse(is_prime(2 * 2      )      )
														self.assertFalse(is_prime(2 * 3      )      )
														self.assertFalse(is_prime(3 * 3      )      )
														self.assertFalse(is_prime(3 * 5      )      )
														self.assertFalse(is_prime(3 * 5 * 7      )      )


if __name__ == "__main__":
				unittest.main()

