



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







							def __init__(   self,			__a,			__a=3,			__a=32,			__a=3,			__a=10,			__a=[10, 20, 30, 40],			__a=[1, 1, 2, 1],			__a=True,			__a=True,			__a="relu",			__a=3,			__a=None,			):

														'''simple docstring'''

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







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Tuple    =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

														_lowerCAmelCase     :							str    =    self.get_config()

														return config, pixel_values







							def    snake_case__      (   self):

														'''simple docstring'''

														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    snake_case__      (   self,			__a,			__a):

														'''simple docstring'''

														_lowerCAmelCase     :							List[Any]    =    FlaxRegNetModel(config=__a)
														_lowerCAmelCase     :							str    =    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 // 32, self.image_size // 32),			)







							def    snake_case__      (   self,			__a,			__a):

														'''simple docstring'''

														_lowerCAmelCase     :							Optional[Any]    =    self.num_labels
														_lowerCAmelCase     :							Tuple    =    FlaxRegNetForImageClassification(config=__a)
														_lowerCAmelCase     :							Union[str, Any]    =    model(__a)
														self.parent.assertEqual(result.logits.shape,			(self.batch_size, self.num_labels))







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Union[str, Any]    =    self.prepare_config_and_inputs()
														_lowerCAmelCase   ,    _lowerCAmelCase     :							int    =    config_and_inputs
														_lowerCAmelCase     :							Optional[int]    =    {"pixel_values": pixel_values}
														return config, inputs_dict



@require_flax
class   UpperCAmelCase_    (       a   ,					unittest.TestCase):
							lowerCamelCase__   			=		(FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

							lowerCamelCase__   			=		False
							lowerCamelCase__   			=		False
							lowerCamelCase__   			=		False







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							List[Any]    =    FlaxRegNetModelTester(self)
														_lowerCAmelCase     :							Optional[Any]    =    ConfigTester(self,			config_class=__a,			has_text_modality=__a)







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

														'''simple docstring'''

														return







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							List[Any]    =    self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*__a)







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Optional[Any]    =    self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_image_classification(*__a)







							@unittest.skip(reason="RegNet does not use inputs_embeds")
							def    snake_case__      (   self):

														'''simple docstring'''

														pass







							@unittest.skip(reason="RegNet does not support input and output embeddings")
							def    snake_case__      (   self):

														'''simple docstring'''

														pass







							def    snake_case__      (   self):

														'''simple docstring'''

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

														for model_class in self.all_model_classes:
																					_lowerCAmelCase     :							Tuple    =    model_class(__a)
																					_lowerCAmelCase     :							Any    =    inspect.signature(model.__call__)
																					# signature.parameters is an OrderedDict => so arg_names order is deterministic
																					_lowerCAmelCase     :							Union[str, Any]    =    [*signature.parameters.keys()]

																					_lowerCAmelCase     :							Dict    =    ["pixel_values"]
																					self.assertListEqual(arg_names[:1],			__a)







							def    snake_case__      (   self):

														'''simple docstring'''

														def check_hidden_states_output(__a,			__a,			__a):
																					_lowerCAmelCase     :							Any    =    model_class(__a)

																					_lowerCAmelCase     :							Tuple    =    model(**self._prepare_for_class(__a,			__a))

																					_lowerCAmelCase     :							Optional[int]    =    outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																					_lowerCAmelCase     :							Dict    =    self.model_tester.num_stages
																					self.assertEqual(len(__a),			expected_num_stages + 1)

														_lowerCAmelCase   ,    _lowerCAmelCase     :							int    =    self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					_lowerCAmelCase     :							Optional[int]    =    True
																					check_hidden_states_output(__a,			__a,			__a)

																					# check that output_hidden_states also work using config
																					del inputs_dict["output_hidden_states"]
																					_lowerCAmelCase     :							Optional[Any]    =    True

																					check_hidden_states_output(__a,			__a,			__a)







							def    snake_case__      (   self):

														'''simple docstring'''

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

														for model_class in self.all_model_classes:
																					with self.subTest(model_class.__name__):
																												_lowerCAmelCase     :							Tuple    =    self._prepare_for_class(__a,			__a)
																												_lowerCAmelCase     :							List[str]    =    model_class(__a)

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

																												with self.subTest("JIT Enabled"):
																																			_lowerCAmelCase     :							Dict    =    model_jitted(**__a).to_tuple()

																												with self.subTest("JIT Disabled"):
																																			with jax.disable_jit():
																																										_lowerCAmelCase     :							Optional[int]    =    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 				A				(					):





							'''simple docstring'''




							_lowerCAmelCase     :							Tuple    =    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"  )
							return image



@require_flax
class   UpperCAmelCase_    (       unittest.TestCase):







							@cached_property
							def    snake_case__      (   self):

														'''simple docstring'''

														return AutoImageProcessor.from_pretrained("facebook/regnet-y-040") if is_vision_available() else None







							@slow
							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    FlaxRegNetForImageClassification.from_pretrained("facebook/regnet-y-040")

														_lowerCAmelCase     :							Dict    =    self.default_image_processor
														_lowerCAmelCase     :							str    =    prepare_img()
														_lowerCAmelCase     :							Tuple    =    image_processor(images=__a,			return_tensors="np")

														_lowerCAmelCase     :							str    =    model(**__a)

														# verify the logits
														_lowerCAmelCase     :							List[Any]    =    (1, 1000)
														self.assertEqual(outputs.logits.shape,			__a)

														_lowerCAmelCase     :							Tuple    =    jnp.array([-0.4_180, -1.5_051, -3.4_836])

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





from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_torch_dist_unique_port,
    require_torch_multi_gpu,
    require_torch_neuroncore,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging


_snake_case		=  logging.get_logger(__name__)


if is_torch_available():
						import torch
						from torch import nn
						from torch.utils.data import Dataset

						from transformers import Trainer



						class   UpperCAmelCase_    (       a):







													def __init__(   self,			__a = 101):

																				'''simple docstring'''

																				_lowerCAmelCase     :							str    =    length







													def __len__(   self):

																				'''simple docstring'''

																				return self.length







													def __getitem__(   self,			__a):

																				'''simple docstring'''

																				return i



						class   UpperCAmelCase_    :







													def __call__(   self,			__a):

																				'''simple docstring'''

																				return {"input_ids": torch.tensor(__a), "labels": torch.tensor(__a)}



						class   UpperCAmelCase_    (       nn.Module):







													def __init__(   self):

																				'''simple docstring'''

																				super().__init__()
																				# Add some (unused) params otherwise DDP will complain.
																				_lowerCAmelCase     :							str    =    nn.Linear(120,			80)







													def    snake_case__      (   self,			__a,			__a=None):

																				'''simple docstring'''

																				if labels is not None:
																											return torch.tensor(0.0,			device=input_ids.device), input_ids
																				else:
																											return input_ids



class   UpperCAmelCase_    (       a):







							@require_torch_neuroncore
							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    f"--nproc_per_node=2\n            --master_port={get_torch_dist_unique_port()}\n            {self.test_file_dir}/test_trainer_distributed.py\n        ".split()
														_lowerCAmelCase     :							Tuple    =    self.get_auto_remove_tmp_dir()
														_lowerCAmelCase     :							Optional[int]    =    f"--output_dir {output_dir}".split()
														_lowerCAmelCase     :							List[Any]    =    ["torchrun"] + distributed_args + args
														execute_subprocess_async(__a,			env=self.get_env())
														# successful return here == success - any errors would have caused an error in the sub-call



class   UpperCAmelCase_    (       a):







							@require_torch_multi_gpu
							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Dict    =    f"--nproc_per_node={torch.cuda.device_count()}\n            --master_port={get_torch_dist_unique_port()}\n            {self.test_file_dir}/test_trainer_distributed.py\n        ".split()
														_lowerCAmelCase     :							Any    =    self.get_auto_remove_tmp_dir()
														_lowerCAmelCase     :							Optional[int]    =    f"--output_dir {output_dir}".split()
														_lowerCAmelCase     :							Any    =    ["torchrun"] + distributed_args + args
														execute_subprocess_async(__a,			env=self.get_env())
														# successful return here == success - any errors would have caused an error in the sub-call


if __name__ == "__main__":
						# The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
						#
						# PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py

						_snake_case		=  HfArgumentParser((TrainingArguments,))
						_snake_case		=  parser.parse_args_into_dataclasses()[0]

						logger.warning(
						    f'''Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, '''
						    f'''distributed training: {training_args.parallel_mode != ParallelMode.NOT_DISTRIBUTED}'''
						)

						# Essentially, what we want to verify in the distributed case is that we get all samples back,
						# in the right order. (this is crucial for prediction for instance)
						for dataset_length in [101, 40, 7]:
												_snake_case		=  DummyDataset(dataset_length)






												def 				A				(					_lowerCamelCase  ):





																			'''simple docstring'''




																			_lowerCAmelCase     :							Dict    =    list(range(len(_lowerCamelCase  )  )  )
																			_lowerCAmelCase     :							Union[str, Any]    =    p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
																			if not success and training_args.local_rank == 0:
																										logger.warning(
																										    "Predictions and/or labels do not match expected results:\n  - predictions: "
																										    F"{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}"  )
																			return {"success": success}

												_snake_case		=  Trainer(
												    model=DummyModel(),
												    args=training_args,
												    data_collator=DummyDataCollator(),
												    eval_dataset=dataset,
												    compute_metrics=compute_metrics,
												)
												_snake_case		=  trainer.evaluate()
												logger.info(metrics)
												if metrics["eval_success"] is not True:
																		logger.error(metrics)
																		exit(1)

												_snake_case		=  trainer.predict(dataset)
												logger.info(p.metrics)
												if p.metrics["test_success"] is not True:
																		logger.error(p.metrics)
																		exit(1)

												_snake_case		=  2

												_snake_case		=  trainer.evaluate()
												logger.info(metrics)
												if metrics["eval_success"] is not True:
																		logger.error(metrics)
																		exit(1)

												_snake_case		=  trainer.predict(dataset)
												logger.info(p.metrics)
												if p.metrics["test_success"] is not True:
																		logger.error(p.metrics)
																		exit(1)



												_snake_case		=  None
