


import inspect
import unittest

from transformers import ResNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

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


if is_flax_available():
				import jax
				import jax.numpy as jnp

				from transformers.models.resnet.modeling_flax_resnet import FlaxResNetForImageClassification, FlaxResNetModel

if is_vision_available():
				from PIL import Image

				from transformers import AutoImageProcessor







class   __lowercase						(unittest.TestCase     ):





							"""simple docstring"""







							def __init__(			self ,	lowerCAmelCase__ ,	lowerCAmelCase__=3 ,	lowerCAmelCase__=3_2 ,	lowerCAmelCase__=3 ,	lowerCAmelCase__=1_0 ,	lowerCAmelCase__=[1_0, 2_0, 3_0, 4_0] ,	lowerCAmelCase__=[1, 1, 2, 1] ,	lowerCAmelCase__=True ,	lowerCAmelCase__=True ,	lowerCAmelCase__="relu" ,	lowerCAmelCase__=3 ,	lowerCAmelCase__=None ,	):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:    List[Any] 		=				parent
										SCREAMING_SNAKE_CASE_		:    str 		=				batch_size
										SCREAMING_SNAKE_CASE_		:    Tuple 		=				image_size
										SCREAMING_SNAKE_CASE_		:    Optional[Any] 		=				num_channels
										SCREAMING_SNAKE_CASE_		:    str 		=				embeddings_size
										SCREAMING_SNAKE_CASE_		:    Tuple 		=				hidden_sizes
										SCREAMING_SNAKE_CASE_		:    Union[str, Any] 		=				depths
										SCREAMING_SNAKE_CASE_		:    List[Any] 		=				is_training
										SCREAMING_SNAKE_CASE_		:    Union[str, Any] 		=				use_labels
										SCREAMING_SNAKE_CASE_		:    Optional[Any] 		=				hidden_act
										SCREAMING_SNAKE_CASE_		:    Union[str, Any] 		=				num_labels
										SCREAMING_SNAKE_CASE_		:    str 		=				scope
										SCREAMING_SNAKE_CASE_		:    List[str] 		=				len(lowerCAmelCase__					)






							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:    Any 		=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

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

										return config, pixel_values






							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										return ResNetConfig(
										    num_channels=self.num_channels ,	embeddings_size=self.embeddings_size ,	hidden_sizes=self.hidden_sizes ,	depths=self.depths ,	hidden_act=self.hidden_act ,	num_labels=self.num_labels ,	image_size=self.image_size ,	)






							def 				UpperCamelCase__    (			self ,	lowerCAmelCase__ ,	lowerCAmelCase__					):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:    Optional[int] 		=				FlaxResNetModel(config=lowerCAmelCase__					)
										SCREAMING_SNAKE_CASE_		:    Optional[Any] 		=				model(lowerCAmelCase__					)

										# Output shape (b, c, h, w)
										self.parent.assertEqual(
										    result.last_hidden_state.shape ,	(self.batch_size, self.hidden_sizes[-1], self.image_size // 3_2, self.image_size // 3_2) ,	)






							def 				UpperCamelCase__    (			self ,	lowerCAmelCase__ ,	lowerCAmelCase__					):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:    List[Any] 		=				self.num_labels
										SCREAMING_SNAKE_CASE_		:    Union[str, Any] 		=				FlaxResNetForImageClassification(config=lowerCAmelCase__					)
										SCREAMING_SNAKE_CASE_		:    Union[str, Any] 		=				model(lowerCAmelCase__					)
										self.parent.assertEqual(result.logits.shape ,	(self.batch_size, self.num_labels)					)






							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:    List[Any] 		=				self.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE_		,			SCREAMING_SNAKE_CASE_		:    Union[str, Any] 		=				config_and_inputs
										SCREAMING_SNAKE_CASE_		:    str 		=				{'pixel_values': pixel_values}
										return config, inputs_dict







@require_flax
class   __lowercase						(__SCREAMING_SNAKE_CASE  ,					unittest.TestCase     ):





							"""simple docstring"""







							_UpperCAmelCase						     =	(FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

							_UpperCAmelCase						     =	False
							_UpperCAmelCase						     =	False
							_UpperCAmelCase						     =	False






							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:    Any 		=				FlaxResNetModelTester(self					)
										SCREAMING_SNAKE_CASE_		:    Dict 		=				ConfigTester(self ,	config_class=lowerCAmelCase__ ,	has_text_modality=lowerCAmelCase__					)






							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										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 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										return






							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:    int 		=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowerCAmelCase__					)






							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:    Optional[int] 		=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase__					)






							@unittest.skip(reason='ResNet does not use inputs_embeds'					)
							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										pass






							@unittest.skip(reason='ResNet does not support input and output embeddings'					)
							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										pass






							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




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

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

													SCREAMING_SNAKE_CASE_		:    Tuple 		=				['pixel_values']
													self.assertListEqual(arg_names[:1] ,	lowerCAmelCase__					)






							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										def check_hidden_states_output(lowerCAmelCase__ ,	lowerCAmelCase__ ,	lowerCAmelCase__					):
													SCREAMING_SNAKE_CASE_		:    Tuple 		=				model_class(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE_		:    int 		=				model(**self._prepare_for_class(lowerCAmelCase__ ,	lowerCAmelCase__					)					)

													SCREAMING_SNAKE_CASE_		:    List[Any] 		=				outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

													SCREAMING_SNAKE_CASE_		:    List[Any] 		=				self.model_tester.num_stages
													self.assertEqual(len(lowerCAmelCase__					) ,	expected_num_stages + 1					)






							@unittest.skip(reason='ResNet does not use feedforward chunking'					)
							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										pass






							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		,			SCREAMING_SNAKE_CASE_		:    Optional[Any] 		=				self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													with self.subTest(model_class.__name__					):
																SCREAMING_SNAKE_CASE_		:    int 		=				self._prepare_for_class(lowerCAmelCase__ ,	lowerCAmelCase__					)
																SCREAMING_SNAKE_CASE_		:    List[Any] 		=				model_class(lowerCAmelCase__					)

																@jax.jit
																def model_jitted(lowerCAmelCase__ ,	**lowerCAmelCase__					):
																			return model(pixel_values=lowerCAmelCase__ ,	**lowerCAmelCase__					)

																with self.subTest('JIT Enabled'					):
																			SCREAMING_SNAKE_CASE_		:    Optional[Any] 		=				model_jitted(**lowerCAmelCase__					).to_tuple()

																with self.subTest('JIT Disabled'					):
																			with jax.disable_jit():
																						SCREAMING_SNAKE_CASE_		:    List[str] 		=				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					)







def 	a__ (							):
			SCREAMING_SNAKE_CASE_		:    List[Any] 		=				Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'							)
			return image







@require_flax
class   __lowercase						(unittest.TestCase     ):





							"""simple docstring"""







							@cached_property
							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										return AutoImageProcessor.from_pretrained('microsoft/resnet-50'					) if is_vision_available() else None






							@slow
							def 				UpperCamelCase__    (			self					):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:    List[str] 		=				FlaxResNetForImageClassification.from_pretrained('microsoft/resnet-50'					)

										SCREAMING_SNAKE_CASE_		:    int 		=				self.default_image_processor
										SCREAMING_SNAKE_CASE_		:    List[Any] 		=				prepare_img()
										SCREAMING_SNAKE_CASE_		:    Union[str, Any] 		=				image_processor(images=lowerCAmelCase__ ,	return_tensors='np'					)

										SCREAMING_SNAKE_CASE_		:    Any 		=				model(**lowerCAmelCase__					)

										# verify the logits
										SCREAMING_SNAKE_CASE_		:    Optional[Any] 		=				(1, 1_0_0_0)
										self.assertEqual(outputs.logits.shape ,	lowerCAmelCase__					)

										SCREAMING_SNAKE_CASE_		:    Any 		=				jnp.array([-11.1_069, -9.7_877, -8.3_777]					)

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









from itertools import compress, repeat
from math import ceil, sqrt






def 	a__ (							A__							):

			if num <= 2:
						return []
			if num == 3:
						return [2]

			# Odd sieve for numbers in range [3, num - 1]
			SCREAMING_SNAKE_CASE_		:    Any 		=				bytearray(B'\x01'							) * ((num >> 1) - 1)

			for i in range(3,			int(sqrt(A__							)							) + 1,			2							):
						if sieve[(i >> 1) - 1]:
									SCREAMING_SNAKE_CASE_		:    Optional[int] 		=				i**2
									SCREAMING_SNAKE_CASE_		:    List[str] 		=				repeat(
									    0,			ceil((num - i_squared) / (i << 1)							)							)

			return [2] + list(compress(range(3,			A__,			2							),			A__							)							)


if __name__ == "__main__":
				import doctest

				doctest.testmod()





