
import inspect
import unittest

from transformers import RegNetConfig, 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.regnet.modeling_flax_regnet import FlaxRegNetForImageClassification, FlaxRegNetModel

if is_vision_available():
					from PIL import Image

					from transformers import AutoImageProcessor

class   lowerCAmelCase__	(      unittest.TestCase     ):






		def __init__(						self						:				Tuple				,     SCREAMING_SNAKE_CASE__						:				str				,     SCREAMING_SNAKE_CASE__						:				str=3				,     SCREAMING_SNAKE_CASE__						:				Tuple=32				,     SCREAMING_SNAKE_CASE__						:				Tuple=3				,     SCREAMING_SNAKE_CASE__						:				Optional[int]=10				,     SCREAMING_SNAKE_CASE__						:				List[Any]=[10, 20, 30, 40]				,     SCREAMING_SNAKE_CASE__						:				Union[str, Any]=[1, 1, 2, 1]				,     SCREAMING_SNAKE_CASE__						:				Optional[int]=True				,     SCREAMING_SNAKE_CASE__						:				List[str]=True				,     SCREAMING_SNAKE_CASE__						:				List[str]="relu"				,     SCREAMING_SNAKE_CASE__						:				Any=3				,     SCREAMING_SNAKE_CASE__						:				Any=None				,     )	->				Any:
						__lowerCamelCase          =      parent
						__lowerCamelCase          =      batch_size
						__lowerCamelCase          =      image_size
						__lowerCamelCase          =      num_channels
						__lowerCamelCase          =      embeddings_size
						__lowerCamelCase          =      hidden_sizes
						__lowerCamelCase          =      depths
						__lowerCamelCase          =      is_training
						__lowerCamelCase          =      use_labels
						__lowerCamelCase          =      hidden_act
						__lowerCamelCase          =      num_labels
						__lowerCamelCase          =      scope
						__lowerCamelCase          =      len(SCREAMING_SNAKE_CASE_    )






		def 					__A			(						self						:				int    )	->				Dict:
						__lowerCamelCase          =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

						__lowerCamelCase          =      self.get_config()

						return config, pixel_values






		def 					__A			(						self						:				Optional[int]    )	->				Optional[int]:
						return RegNetConfig(
						    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 					__A			(						self						:				Any				,     SCREAMING_SNAKE_CASE__						:				Dict				,     SCREAMING_SNAKE_CASE__						:				Dict    )	->				List[str]:
						__lowerCamelCase          =      FlaxRegNetModel(config=SCREAMING_SNAKE_CASE_    )
						__lowerCamelCase          =      model(SCREAMING_SNAKE_CASE_    )

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






		def 					__A			(						self						:				Any				,     SCREAMING_SNAKE_CASE__						:				str				,     SCREAMING_SNAKE_CASE__						:				str    )	->				str:
						__lowerCamelCase          =      self.num_labels
						__lowerCamelCase          =      FlaxRegNetForImageClassification(config=SCREAMING_SNAKE_CASE_    )
						__lowerCamelCase          =      model(SCREAMING_SNAKE_CASE_    )
						self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.num_labels)    )







		def 					__A			(						self						:				Any    )	->				Optional[int]:
						__lowerCamelCase          =      self.prepare_config_and_inputs()
						__lowerCamelCase   ,       __lowerCamelCase          =      config_and_inputs
						__lowerCamelCase          =      {'''pixel_values''': pixel_values}
						return config, inputs_dict





@require_flax
class   lowerCAmelCase__	(      lowercase__    ,      unittest.TestCase     ):
		a__			:		Optional[Any]     =						(FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

		a__			:		int     =						False
		a__			:		Union[str, Any]     =						False
		a__			:		Dict     =						False






		def 					__A			(						self						:				int    )	->				None:
						__lowerCamelCase          =      FlaxRegNetModelTester(self    )
						__lowerCamelCase          =      ConfigTester(self				,     config_class=SCREAMING_SNAKE_CASE_				,     has_text_modality=SCREAMING_SNAKE_CASE_    )






		def 					__A			(						self						:				Any    )	->				Any:
						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 					__A			(						self						:				Tuple    )	->				Tuple:
						return






		def 					__A			(						self						:				int    )	->				str:
						__lowerCamelCase          =      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_    )






		def 					__A			(						self						:				Tuple    )	->				str:
						__lowerCamelCase          =      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE_    )






		@unittest.skip(reason='''RegNet does not use inputs_embeds'''    )
		def 					__A			(						self						:				Optional[Any]    )	->				Tuple:
						pass






		@unittest.skip(reason='''RegNet does not support input and output embeddings'''    )
		def 					__A			(						self						:				Optional[int]    )	->				Any:
						pass






		def 					__A			(						self						:				Union[str, Any]    )	->				Any:
						__lowerCamelCase   ,       __lowerCamelCase          =      self.model_tester.prepare_config_and_inputs_for_common()

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

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






		def 					__A			(						self						:				Optional[int]    )	->				Dict:


						def check_hidden_states_output(SCREAMING_SNAKE_CASE__						:				Optional[int]				,     SCREAMING_SNAKE_CASE__						:				List[Any]				,     SCREAMING_SNAKE_CASE__						:				Union[str, Any]    ):
										__lowerCamelCase          =      model_class(SCREAMING_SNAKE_CASE_    )

										__lowerCamelCase          =      model(**self._prepare_for_class(SCREAMING_SNAKE_CASE_				,     SCREAMING_SNAKE_CASE_    )    )

										__lowerCamelCase          =      outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

										__lowerCamelCase          =      self.model_tester.num_stages
										self.assertEqual(len(SCREAMING_SNAKE_CASE_    )				,     expected_num_stages + 1    )

						__lowerCamelCase   ,       __lowerCamelCase          =      self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										__lowerCamelCase          =      True
										check_hidden_states_output(SCREAMING_SNAKE_CASE_				,     SCREAMING_SNAKE_CASE_				,     SCREAMING_SNAKE_CASE_    )

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

										check_hidden_states_output(SCREAMING_SNAKE_CASE_				,     SCREAMING_SNAKE_CASE_				,     SCREAMING_SNAKE_CASE_    )







		def 					__A			(						self						:				List[Any]    )	->				List[str]:
						__lowerCamelCase   ,       __lowerCamelCase          =      self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										with self.subTest(model_class.__name__    ):
														__lowerCamelCase          =      self._prepare_for_class(SCREAMING_SNAKE_CASE_				,     SCREAMING_SNAKE_CASE_    )
														__lowerCamelCase          =      model_class(SCREAMING_SNAKE_CASE_    )

														@jax.jit
														def model_jitted(SCREAMING_SNAKE_CASE__						:				Optional[Any]				,     **SCREAMING_SNAKE_CASE__						:				Any    ):
																		return model(pixel_values=SCREAMING_SNAKE_CASE_				,     **SCREAMING_SNAKE_CASE_    )

														with self.subTest('''JIT Enabled'''    ):
																		__lowerCamelCase          =      model_jitted(**SCREAMING_SNAKE_CASE_    ).to_tuple()

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

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





def 					__magic_name__					(  )     ->  Tuple:
				__lowerCamelCase          =      Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
				return image







@require_flax
class   lowerCAmelCase__	(      unittest.TestCase     ):






		@cached_property
		def 					__A			(						self						:				str    )	->				Optional[int]:
						return AutoImageProcessor.from_pretrained('''facebook/regnet-y-040'''    ) if is_vision_available() else None







		@slow
		def 					__A			(						self						:				Optional[int]    )	->				Dict:
						__lowerCamelCase          =      FlaxRegNetForImageClassification.from_pretrained('''facebook/regnet-y-040'''    )

						__lowerCamelCase          =      self.default_image_processor
						__lowerCamelCase          =      prepare_img()
						__lowerCamelCase          =      image_processor(images=SCREAMING_SNAKE_CASE_				,     return_tensors='''np'''    )

						__lowerCamelCase          =      model(**SCREAMING_SNAKE_CASE_    )

						# verify the logits
						__lowerCamelCase          =      (1, 10_00)
						self.assertEqual(outputs.logits.shape				,     SCREAMING_SNAKE_CASE_    )

						__lowerCamelCase          =      jnp.array([-0.4180, -1.5051, -3.4836]    )

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





"""simple docstring"""






from ..utils import DummyObject, requires_backends



class 			lowercase					( metaclass=lowercase__		):
			lowercase        =     ['''flax''', '''transformers''']






			def __init__(self  :	List[Any]      ,*SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,**SCREAMING_SNAKE_CASE_  :	Union[str, Any]			)       ->	str:





				"""simple docstring"""


				requires_backends(self      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	List[Any]      ,*SCREAMING_SNAKE_CASE_  :	List[Any]      ,**SCREAMING_SNAKE_CASE_  :	int			)       ->	Optional[int]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	List[Any]      ,*SCREAMING_SNAKE_CASE_  :	str      ,**SCREAMING_SNAKE_CASE_  :	Union[str, Any]			)       ->	Union[str, Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)





class 			lowercase					( metaclass=lowercase__		):
			lowercase        =     ['''flax''', '''transformers''']






			def __init__(self  :	int      ,*SCREAMING_SNAKE_CASE_  :	List[Any]      ,**SCREAMING_SNAKE_CASE_  :	Any			)       ->	Union[str, Any]:





				"""simple docstring"""


				requires_backends(self      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	int      ,*SCREAMING_SNAKE_CASE_  :	str      ,**SCREAMING_SNAKE_CASE_  :	int			)       ->	Optional[Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	List[str]      ,*SCREAMING_SNAKE_CASE_  :	List[str]      ,**SCREAMING_SNAKE_CASE_  :	Optional[Any]			)       ->	Dict:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)





class 			lowercase					( metaclass=lowercase__		):
			lowercase        =     ['''flax''', '''transformers''']






			def __init__(self  :	Tuple      ,*SCREAMING_SNAKE_CASE_  :	Optional[int]      ,**SCREAMING_SNAKE_CASE_  :	Dict			)       ->	List[str]:





				"""simple docstring"""


				requires_backends(self      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	Tuple      ,*SCREAMING_SNAKE_CASE_  :	Dict      ,**SCREAMING_SNAKE_CASE_  :	Optional[Any]			)       ->	List[Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	Tuple      ,*SCREAMING_SNAKE_CASE_  :	Tuple      ,**SCREAMING_SNAKE_CASE_  :	Tuple			)       ->	Union[str, Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)





class 			lowercase					( metaclass=lowercase__		):
			lowercase        =     ['''flax''', '''transformers''']






			def __init__(self  :	Optional[Any]      ,*SCREAMING_SNAKE_CASE_  :	Any      ,**SCREAMING_SNAKE_CASE_  :	str			)       ->	Union[str, Any]:





				"""simple docstring"""


				requires_backends(self      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	List[Any]      ,*SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,**SCREAMING_SNAKE_CASE_  :	Optional[int]			)       ->	Union[str, Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	List[Any]      ,*SCREAMING_SNAKE_CASE_  :	List[Any]      ,**SCREAMING_SNAKE_CASE_  :	List[Any]			)       ->	Optional[Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)


