


'''simple docstring'''


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 						lowerCAmelCase__      (      unittest.TestCase						):




			def __init__(		self	:  Any						,	lowerCamelCase__	:  int						,	lowerCamelCase__	:  List[str]=3						,	lowerCamelCase__	:  List[Any]=32						,	lowerCamelCase__	:  Union[str, Any]=3						,	lowerCamelCase__	:  Union[str, Any]=10						,	lowerCamelCase__	:  List[Any]=[10, 20, 30, 40]						,	lowerCamelCase__	:  str=[1, 1, 2, 1]						,	lowerCamelCase__	:  List[str]=True						,	lowerCamelCase__	:  int=True						,	lowerCamelCase__	:  Dict="relu"						,	lowerCamelCase__	:  str=3						,	lowerCamelCase__	:  Optional[int]=None						,	) ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase		:							str             =   parent
					_UpperCAmelCase		:							List[Any]             =   batch_size
					_UpperCAmelCase		:							Dict             =   image_size
					_UpperCAmelCase		:							str             =   num_channels
					_UpperCAmelCase		:							Optional[Any]             =   embeddings_size
					_UpperCAmelCase		:							str             =   hidden_sizes
					_UpperCAmelCase		:							Optional[Any]             =   depths
					_UpperCAmelCase		:							str             =   is_training
					_UpperCAmelCase		:							Dict             =   use_labels
					_UpperCAmelCase		:							Any             =   hidden_act
					_UpperCAmelCase		:							Tuple             =   num_labels
					_UpperCAmelCase		:							Optional[Any]             =   scope
					_UpperCAmelCase		:							str             =   len(lowerCamelCase__			)




			def 				lowerCAmelCase__     (		self	:  Optional[int]			) ->int:



					'''simple docstring'''







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

					_UpperCAmelCase		:							Dict             =   self.get_config()

					return config, pixel_values




			def 				lowerCAmelCase__     (		self	:  Dict			) ->Tuple:



					'''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 				lowerCAmelCase__     (		self	:  Tuple						,	lowerCamelCase__	:  Dict						,	lowerCamelCase__	:  str			) ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase		:							Union[str, Any]             =   FlaxResNetModel(config=lowerCamelCase__			)
					_UpperCAmelCase		:							str             =   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 // 32, self.image_size // 32)						,	)




			def 				lowerCAmelCase__     (		self	:  List[str]						,	lowerCamelCase__	:  str						,	lowerCamelCase__	:  Optional[int]			) ->Any:



					'''simple docstring'''







					_UpperCAmelCase		:							Optional[int]             =   self.num_labels
					_UpperCAmelCase		:							Optional[Any]             =   FlaxResNetForImageClassification(config=lowerCamelCase__			)
					_UpperCAmelCase		:							List[str]             =   model(lowerCamelCase__			)
					self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.num_labels)			)






			def 				lowerCAmelCase__     (		self	:  Tuple			) ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							List[Any]             =   self.prepare_config_and_inputs()
					_UpperCAmelCase       ,     _UpperCAmelCase		:							Tuple             =   config_and_inputs
					_UpperCAmelCase		:							Optional[int]             =   {"pixel_values": pixel_values}
					return config, inputs_dict






@require_flax
class 						lowerCAmelCase__      (      UpperCAmelCase__     ,							unittest.TestCase						):
			lowerCAmelCase			:						Tuple							 =  (FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

			lowerCAmelCase			:						Dict							 =  False
			lowerCAmelCase			:						Any							 =  False
			lowerCAmelCase			:						Optional[Any]							 =  False




			def 				lowerCAmelCase__     (		self	:  Optional[Any]			) ->None:



					'''simple docstring'''







					_UpperCAmelCase		:							List[Any]             =   FlaxResNetModelTester(self			)
					_UpperCAmelCase		:							Union[str, Any]             =   ConfigTester(self						,	config_class=lowerCamelCase__						,	has_text_modality=lowerCamelCase__			)




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



					'''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 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->Tuple:



					'''simple docstring'''







					return




			def 				lowerCAmelCase__     (		self	:  str			) ->int:



					'''simple docstring'''







					_UpperCAmelCase		:							List[Any]             =   self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*lowerCamelCase__			)




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



					'''simple docstring'''







					_UpperCAmelCase		:							Union[str, Any]             =   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 				lowerCAmelCase__     (		self	:  Union[str, Any]			) ->List[Any]:



					'''simple docstring'''







					pass




			@unittest.skip(reason="ResNet does not support input and output embeddings"			)
			def 				lowerCAmelCase__     (		self	:  Dict			) ->str:



					'''simple docstring'''







					pass




			def 				lowerCAmelCase__     (		self	:  Any			) ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Dict             =   self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							_UpperCAmelCase		:							Optional[Any]             =   model_class(lowerCamelCase__			)
							_UpperCAmelCase		:							str             =   inspect.signature(model.__call__			)
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							_UpperCAmelCase		:							Optional[int]             =   [*signature.parameters.keys()]

							_UpperCAmelCase		:							List[str]             =   ["pixel_values"]
							self.assertListEqual(arg_names[:1]						,	lowerCamelCase__			)




			def 				lowerCAmelCase__     (		self	:  str			) ->Tuple:



					'''simple docstring'''







					def check_hidden_states_output(lowerCamelCase__	:  Union[str, Any]						,	lowerCamelCase__	:  str						,	lowerCamelCase__	:  Optional[int]			):
							_UpperCAmelCase		:							Tuple             =   model_class(lowerCamelCase__			)

							_UpperCAmelCase		:							str             =   model(**self._prepare_for_class(lowerCamelCase__						,	lowerCamelCase__			)			)

							_UpperCAmelCase		:							Optional[Any]             =   outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

							_UpperCAmelCase		:							List[str]             =   self.model_tester.num_stages
							self.assertEqual(len(lowerCamelCase__			)						,	expected_num_stages + 1			)




			@unittest.skip(reason="ResNet does not use feedforward chunking"			)
			def 				lowerCAmelCase__     (		self	:  int			) ->int:



					'''simple docstring'''







					pass






			def 				lowerCAmelCase__     (		self	:  List[str]			) ->Dict:



					'''simple docstring'''







					_UpperCAmelCase       ,     _UpperCAmelCase		:							Optional[int]             =   self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							with self.subTest(model_class.__name__			):
									_UpperCAmelCase		:							str             =   self._prepare_for_class(lowerCamelCase__						,	lowerCamelCase__			)
									_UpperCAmelCase		:							List[Any]             =   model_class(lowerCamelCase__			)

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

									with self.subTest("JIT Enabled"			):
											_UpperCAmelCase		:							Optional[int]             =   model_jitted(**lowerCamelCase__			).to_tuple()

									with self.subTest("JIT Disabled"			):
											with jax.disable_jit():
													_UpperCAmelCase		:							Optional[Any]             =   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 			__lowerCAmelCase			():
		_UpperCAmelCase		:							Optional[Any]             =   Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png" )
		return image






@require_flax
class 						lowerCAmelCase__      (      unittest.TestCase						):




			@cached_property
			def 				lowerCAmelCase__     (		self	:  str			) ->int:



					'''simple docstring'''







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






			@slow
			def 				lowerCAmelCase__     (		self	:  int			) ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase		:							Optional[int]             =   FlaxResNetForImageClassification.from_pretrained("microsoft/resnet-50"			)

					_UpperCAmelCase		:							List[str]             =   self.default_image_processor
					_UpperCAmelCase		:							Dict             =   prepare_img()
					_UpperCAmelCase		:							int             =   image_processor(images=lowerCamelCase__						,	return_tensors="np"			)

					_UpperCAmelCase		:							Tuple             =   model(**lowerCamelCase__			)

					# verify the logits
					_UpperCAmelCase		:							Optional[Any]             =   (1, 10_00)
					self.assertEqual(outputs.logits.shape						,	lowerCamelCase__			)

					_UpperCAmelCase		:							Any             =   jnp.array([-1_1.1_0_6_9, -9.7_8_7_7, -8.3_7_7_7]			)

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







'''simple docstring'''


import torch
from torch import autograd






class 						lowerCAmelCase__      (      autograd.Function						):




			@staticmethod
			def 				lowerCAmelCase__     (		lowerCamelCase__	:  Optional[int]						,	lowerCamelCase__	:  torch.tensor						,	lowerCamelCase__	:  float						,	lowerCamelCase__	:  bool			) ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							Optional[int]             =   inputs.numel()
					_UpperCAmelCase		:							Any             =   int(0.0_0_5 * nb_elems			) + 1
					if sigmoid:
							_UpperCAmelCase		:							List[str]             =   (torch.sigmoid(lowerCamelCase__			) > threshold).type(inputs.type()			)
					else:
							_UpperCAmelCase		:							List[Any]             =   (inputs > threshold).type(inputs.type()			)
					if mask.sum() < nb_min:
							# We limit the pruning so that at least 0.5% (half a percent) of the weights are remaining
							_UpperCAmelCase		:							str             =   inputs.flatten().kthvalue(max(nb_elems - nb_min						,	1			)			).values
							_UpperCAmelCase		:							str             =   (inputs > k_threshold).type(inputs.type()			)
					return mask






			@staticmethod
			def 				lowerCAmelCase__     (		lowerCamelCase__	:  List[Any]						,	lowerCamelCase__	:  Optional[int]			) ->Dict:



					'''simple docstring'''







					return gradOutput, None, None






class 						lowerCAmelCase__      (      autograd.Function						):




			@staticmethod
			def 				lowerCAmelCase__     (		lowerCamelCase__	:  Optional[int]						,	lowerCamelCase__	:  torch.tensor						,	lowerCamelCase__	:  float			) ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase		:							Optional[int]             =   inputs.clone()
					_UpperCAmelCase       ,     _UpperCAmelCase		:							Tuple             =   inputs.flatten().sort(descending=lowerCamelCase__			)
					_UpperCAmelCase		:							Optional[int]             =   int(threshold * inputs.numel()			)

					# flat_out and mask access the same memory.
					_UpperCAmelCase		:							str             =   mask.flatten()
					_UpperCAmelCase		:							Optional[int]             =   0
					_UpperCAmelCase		:							Optional[Any]             =   1
					return mask






			@staticmethod
			def 				lowerCAmelCase__     (		lowerCamelCase__	:  Optional[int]						,	lowerCamelCase__	:  Any			) ->Union[str, Any]:



					'''simple docstring'''







					return gradOutput, None






class 						lowerCAmelCase__      (      UpperCAmelCase__						):




			@staticmethod
			def 				lowerCAmelCase__     (		lowerCamelCase__	:  torch.tensor						,	lowerCamelCase__	:  float			) ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase		:							List[Any]             =   inputs.clone()
					_UpperCAmelCase       ,     _UpperCAmelCase		:							Optional[Any]             =   inputs.abs().flatten().sort(descending=lowerCamelCase__			)
					_UpperCAmelCase		:							int             =   int(threshold * inputs.numel()			)

					# flat_out and mask access the same memory.
					_UpperCAmelCase		:							int             =   mask.flatten()
					_UpperCAmelCase		:							List[Any]             =   0
					_UpperCAmelCase		:							int             =   1
					return mask



