

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


							"""simple docstring"""





							def __init__(			self       ,    _A       ,    _A=3       ,    _A=3_2       ,    _A=3       ,    _A=1_0       ,    _A=[1_0, 2_0, 3_0, 4_0]       ,    _A=[1, 1, 2, 1]       ,    _A=True       ,    _A=True       ,    _A="relu"       ,    _A=3       ,    _A=None       ,    ):
												__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(_A				)




							def 							A__	(			self				):
												__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				):
												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 							A__	(			self       ,    _A       ,    _A				):
												__lowerCAmelCase       =				FlaxResNetModel(config=_A				)
												__lowerCAmelCase       =				model(_A				)

												# 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 							A__	(			self       ,    _A       ,    _A				):
												__lowerCAmelCase       =				self.num_labels
												__lowerCAmelCase       =				FlaxResNetForImageClassification(config=_A				)
												__lowerCAmelCase       =				model(_A				)
												self.parent.assertEqual(result.logits.shape       ,    (self.batch_size, self.num_labels)				)




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





@require_flax
class      snake_case_      (					lowerCAmelCase_      ,			unittest.TestCase				):


							"""simple docstring"""





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

							__UpperCAmelCase      =False
							__UpperCAmelCase      =False
							__UpperCAmelCase      =False




							def 							A__	(			self				):
												__lowerCAmelCase       =				FlaxResNetModelTester(self				)
												__lowerCAmelCase       =				ConfigTester(self       ,    config_class=_A       ,    has_text_modality=_A				)




							def 							A__	(			self				):
												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				):
												return




							def 							A__	(			self				):
												__lowerCAmelCase       =				self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_A				)




							def 							A__	(			self				):
												__lowerCAmelCase       =				self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_image_classification(*_A				)




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




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




							def 							A__	(			self				):
												__lowerCAmelCase       =				self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																	__lowerCAmelCase       =				model_class(_A				)
																	__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]       ,    _A				)




							def 							A__	(			self				):
												def check_hidden_states_output(_A       ,    _A       ,    _A				):
																	__lowerCAmelCase       =				model_class(_A				)

																	__lowerCAmelCase       =				model(**self._prepare_for_class(_A       ,    _A				)				)

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

																	__lowerCAmelCase       =				self.model_tester.num_stages
																	self.assertEqual(len(_A				)       ,    expected_num_stages + 1				)




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




							def 							A__	(			self				):
												__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(_A       ,    _A				)
																						__lowerCAmelCase       =				model_class(_A				)

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

																						with self.subTest('JIT Enabled'				):
																											__lowerCAmelCase       =				model_jitted(**_A				).to_tuple()

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

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


def 							__lowercase     (	):

					"""simple docstring"""






					__lowerCAmelCase       =				Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
					return image





@require_flax
class      snake_case_      (					unittest.TestCase				):


							"""simple docstring"""





							@cached_property
							def 							A__	(			self				):
												return AutoImageProcessor.from_pretrained('microsoft/resnet-50'				) if is_vision_available() else None




							@slow
							def 							A__	(			self				):
												__lowerCAmelCase       =				FlaxResNetForImageClassification.from_pretrained('microsoft/resnet-50'				)

												__lowerCAmelCase       =				self.default_image_processor
												__lowerCAmelCase       =				prepare_img()
												__lowerCAmelCase       =				image_processor(images=_A       ,    return_tensors='np'				)

												__lowerCAmelCase       =				model(**_A				)

												# verify the logits
												__lowerCAmelCase       =				(1, 1_0_0_0)
												self.assertEqual(outputs.logits.shape       ,    _A				)

												__lowerCAmelCase       =				jnp.array([-11.1069, -9.7877, -8.3777]				)

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











"""simple docstring"""






import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
		import torch

		from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
		import pyanvml.pyanvml as nvml


_lowercase					:					Optional[Any] 		=		logging.get_logger(__name__)

class 			__SCREAMING_SNAKE_CASE      ( lowerCAmelCase_  ):


					'''simple docstring'''





					_a       =				42
					_a       =				42
					_a       =				"PyTorch"



					@property
					def 						snake_case			(			self				:    Dict	)->   List[str]:
								return torch.__version__



					def 						snake_case			(			self				:    Optional[Any],    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   float:
								lowerCamelCase__						:						List[Any]						      =self._prepare_inference_func(lowerCamelCase,    lowerCamelCase,    lowerCamelCase	)
								return self._measure_speed(_inference	)



					def 						snake_case			(			self				:    List[str],    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   [Memory, Optional[MemorySummary]]:
								lowerCamelCase__						:						List[Any]						      =self._prepare_inference_func(lowerCamelCase,    lowerCamelCase,    lowerCamelCase	)
								return self._measure_memory(_inference	)



					def 						snake_case			(			self				:    Any,    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   float:
								lowerCamelCase__						:						int						      =self._prepare_train_func(lowerCamelCase,    lowerCamelCase,    lowerCamelCase	)
								return self._measure_speed(_train	)



					def 						snake_case			(			self				:    List[str],    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   [Memory, Optional[MemorySummary]]:
								lowerCamelCase__						:						str						      =self._prepare_train_func(lowerCamelCase,    lowerCamelCase,    lowerCamelCase	)
								return self._measure_memory(_train	)



					def 						snake_case			(			self				:    List[str],    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   Callable[[], None]:
								lowerCamelCase__						:						List[str]						      =self.config_dict[model_name]

								if self.args.torchscript:
											lowerCamelCase__						:						int						      =True

								lowerCamelCase__						:						Optional[Any]						      =(
								    hasattr(lowerCamelCase,    '''architectures'''	)
								    and isinstance(config.architectures,    lowerCamelCase	)
								    and len(config.architectures	) > 0
								)
								if not self.args.only_pretrain_model and has_model_class_in_config:
											try:
														lowerCamelCase__						:						int						      =config.architectures[0]
														lowerCamelCase__						:						Dict						      =__import__('''transformers''',    fromlist=[model_class]	)
														lowerCamelCase__						:						str						      =getattr(lowerCamelCase,    lowerCamelCase	)
														lowerCamelCase__						:						Tuple						      =model_cls(lowerCamelCase	)
											except ImportError:
														raise ImportError(
														    F'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
														    ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''	)
								else:
											lowerCamelCase__						:						List[str]						      =MODEL_MAPPING[config.__class__](lowerCamelCase	)

								model.eval()
								model.to(self.args.device	)

								# encoder-decoder has vocab size saved differently
								lowerCamelCase__						:						Optional[Any]						      =config.vocab_size if hasattr(lowerCamelCase,    '''vocab_size'''	) else config.encoder.vocab_size
								lowerCamelCase__						:						int						      =torch.randint(lowerCamelCase,    (batch_size, sequence_length),    dtype=torch.long,    device=self.args.device	)

								if self.args.fpaa:
											logger.info('''Running training in Mixed Precision...'''	)
											if not self.args.is_gpu:
														raise ValueError('''Mixed precision is possible only for GPU.'''	)
											# amp seems to have memory leaks so that memory usage
											# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
											model.half()

								if self.args.torchscript:
											with torch.no_grad():
														lowerCamelCase__						:						Optional[Any]						      =torch.jit.trace(lowerCamelCase,    lowerCamelCase	)
								else:
											lowerCamelCase__						:						Tuple						      =model

								def encoder_decoder_forward():
											with torch.no_grad():
														lowerCamelCase__						:						Optional[Any]						      =inference_model(lowerCamelCase,    decoder_input_ids=lowerCamelCase	)
											return outputs

								def encoder_forward():
											with torch.no_grad():
														lowerCamelCase__						:						Any						      =inference_model(lowerCamelCase	)
											return outputs

								lowerCamelCase__						:						Optional[int]						      =encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
								return _forward



					def 						snake_case			(			self				:    Union[str, Any],    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   Callable[[], None]:
								lowerCamelCase__						:						Any						      =self.config_dict[model_name]

								lowerCamelCase__						:						Tuple						      =(
								    hasattr(lowerCamelCase,    '''architectures'''	)
								    and isinstance(config.architectures,    lowerCamelCase	)
								    and len(config.architectures	) > 0
								)
								if not self.args.only_pretrain_model and has_model_class_in_config:
											try:
														lowerCamelCase__						:						List[Any]						      =config.architectures[0]
														lowerCamelCase__						:						int						      =__import__('''transformers''',    fromlist=[model_class]	)
														lowerCamelCase__						:						int						      =getattr(lowerCamelCase,    lowerCamelCase	)
														lowerCamelCase__						:						Dict						      =model_cls(lowerCamelCase	)
											except ImportError:
														raise ImportError(
														    F'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
														    ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'''	)
								else:
											lowerCamelCase__						:						List[str]						      =MODEL_WITH_LM_HEAD_MAPPING[config.__class__](lowerCamelCase	)

								if self.args.torchscript:
											raise NotImplementedError('''Training for torchscript is currently not implemented'''	)
								else:
											lowerCamelCase__						:						str						      =model

								model.train()
								model.to(self.args.device	)

								# encoder-decoder has vocab size saved differently
								lowerCamelCase__						:						Dict						      =config.vocab_size if hasattr(lowerCamelCase,    '''vocab_size'''	) else config.encoder.vocab_size
								lowerCamelCase__						:						int						      =torch.randint(lowerCamelCase,    (batch_size, sequence_length),    dtype=torch.long,    device=self.args.device	)

								if self.args.fpaa:
											logger.info('''Running training in Mixed Precision...'''	)
											if not self.args.is_gpu:
														raise ValueError('''Mixed precision is possible only for GPU.'''	)

											# amp seems to have memory leaks so that memory usage
											# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
											model.half()

								def compute_loss_and_backprob_encoder():
											lowerCamelCase__						:						Optional[int]						      =train_model(lowerCamelCase,    labels=lowerCamelCase	)[0]
											loss.backward()
											return loss

								def compute_loss_and_backprob_encoder_decoder():
											lowerCamelCase__						:						Dict						      =train_model(lowerCamelCase,    decoder_input_ids=lowerCamelCase,    labels=lowerCamelCase	)[0]
											loss.backward()
											return loss

								lowerCamelCase__						:						Optional[Any]						      =(
								    compute_loss_and_backprob_encoder_decoder
								    if config.is_encoder_decoder
								    else compute_loss_and_backprob_encoder
								)
								return _train



					def 						snake_case			(			self				:    Optional[int],    lowerCamelCase				:    Optional[int]	)->   float:
								try:
											if self.args.is_tpu or self.args.torchscript:
														# run additional 10 times to stabilize compilation for tpu and torchscript
														logger.info('''Do inference on TPU or torchscript. Running model 5 times to stabilize compilation'''	)
														timeit.repeat(
														    lowerCamelCase,    repeat=1,    number=5,    )

											# as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
											lowerCamelCase__						:						List[str]						      =timeit.repeat(
											    lowerCamelCase,    repeat=self.args.repeat,    number=10,    )

											if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
														import torch_xla.debug.metrics as met

														self.print_fn(met.metrics_report()	)

											return min(lowerCamelCase	) / 10.0
								except RuntimeError as e:
											self.print_fn(F'''Doesn\'t fit on GPU. {e}'''	)
											return "N/A"




					def 						snake_case			(			self				:    Optional[int],    lowerCamelCase				:    Callable[[], None]	)->   [Memory, MemorySummary]:
								try:
											if self.args.trace_memory_line_by_line:
														lowerCamelCase__						:						Optional[int]						      =start_memory_tracing('''transformers'''	)

											if self.args.is_tpu:
														# tpu
														raise NotImplementedError(
														    '''Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'''
														    ''' `--no-memory` or `args.memory=False`'''	)
											elif self.args.is_gpu:
														if not is_pyanvml_available():
																	logger.warning(
																	    '''py3nvml not installed, we won\'t log GPU memory usage. '''
																	    '''Install py3nvml (pip install py3nvml) to log information about GPU.'''	)
																	lowerCamelCase__						:						str						      ='''N/A'''
														else:
																	logger.info(
																	    '''Measuring total GPU usage on GPU device. Make sure to not have additional processes running'''
																	    ''' on the same GPU.'''	)
																	# init nvml
																	nvml.nvmlInit()
																	func()
																	lowerCamelCase__						:						int						      =nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
																	lowerCamelCase__						:						Optional[Any]						      =nvml.nvmlDeviceGetMemoryInfo(lowerCamelCase	)
																	lowerCamelCase__						:						str						      =meminfo.used
																	lowerCamelCase__						:						List[str]						      =Memory(lowerCamelCase	)
																	# shutdown nvml
																	nvml.nvmlShutdown()
											else:
														# cpu
														lowerCamelCase__						:						List[Any]						      =measure_peak_memory_cpu(lowerCamelCase	)
														lowerCamelCase__						:						List[Any]						      =Memory(lowerCamelCase	) if isinstance(lowerCamelCase,    lowerCamelCase	) else memory_bytes

											if self.args.trace_memory_line_by_line:
														lowerCamelCase__						:						List[str]						      =stop_memory_tracing(lowerCamelCase	)
											else:
														lowerCamelCase__						:						Optional[int]						      =None

											return memory, summary
								except RuntimeError as e:
											self.print_fn(F'''Doesn\'t fit on GPU. {e}'''	)
											return "N/A", None




