






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


   '''simple docstring'''

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

     '''simple docstring'''
     __A						      =			parent
     __A						      =			batch_size
     __A						      =			image_size
     __A						      =			num_channels
     __A						      =			embeddings_size
     __A						      =			hidden_sizes
     __A						      =			depths
     __A						      =			is_training
     __A						      =			use_labels
     __A						      =			hidden_act
     __A						      =			num_labels
     __A						      =			scope
     __A						      =			len(_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[str]  ):

     '''simple docstring'''
     __A						      =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

     __A						      =			self.get_config()

     return config, pixel_values


   def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[int]  ):

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

     '''simple docstring'''
     __A						      =			FlaxResNetModel(config=_lowerCamelCase  )
     __A						      =			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     _SCREAMING_SNAKE_CASE    (     self			:					Any,	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					int  ):

     '''simple docstring'''
     __A						      =			self.num_labels
     __A						      =			FlaxResNetForImageClassification(config=_lowerCamelCase  )
     __A						      =			model(_lowerCamelCase  )
     self.parent.assertEqual(result.logits.shape,	(self.batch_size, self.num_labels)  )






   def     _SCREAMING_SNAKE_CASE    (     self			:					List[Any]  ):

     '''simple docstring'''
     __A						      =			self.prepare_config_and_inputs()
     __A      ,   __A						      =			config_and_inputs
     __A						      =			{'''pixel_values''': pixel_values}
     return config, inputs_dict





@require_flax
class  snake_case		(							_lowerCAmelCase       ,     unittest.TestCase						):


   '''simple docstring'''

   A_						:     List[str]									=			(FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

   A_						:     str									=			False
   A_						:     Dict									=			False
   A_						:     Tuple									=			False


   def     _SCREAMING_SNAKE_CASE    (     self			:					int  ):

     '''simple docstring'''
     __A						      =			FlaxResNetModelTester(self  )
     __A						      =			ConfigTester(self,	config_class=_lowerCamelCase,	has_text_modality=_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Union[str, Any]  ):

     '''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     _SCREAMING_SNAKE_CASE    (     self			:					Optional[int]  ):

     '''simple docstring'''
     return


   def     _SCREAMING_SNAKE_CASE    (     self			:					int  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[Any]  ):

     '''simple docstring'''
     __A						      =			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     _SCREAMING_SNAKE_CASE    (     self			:					str  ):

     '''simple docstring'''
     pass


   @unittest.skip(reason='''ResNet does not support input and output embeddings'''  )
   def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[int]  ):

     '''simple docstring'''
     pass


   def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[int]  ):

     '''simple docstring'''
     __A      ,   __A						      =			self.model_tester.prepare_config_and_inputs_for_common()

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

       __A						      =			['''pixel_values''']
       self.assertListEqual(arg_names[:1],	_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Union[str, Any]  ):

     '''simple docstring'''
     def check_hidden_states_output(_lowerCamelCase			:					Dict,	_lowerCamelCase			:					List[str],	_lowerCamelCase			:					Any  ):
       __A						      =			model_class(_lowerCamelCase  )

       __A						      =			model(**self._prepare_for_class(_lowerCamelCase,	_lowerCamelCase  )  )

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

       __A						      =			self.model_tester.num_stages
       self.assertEqual(len(_lowerCamelCase  ),	expected_num_stages + 1  )


   @unittest.skip(reason='''ResNet does not use feedforward chunking'''  )
   def     _SCREAMING_SNAKE_CASE    (     self			:					List[Any]  ):

     '''simple docstring'''
     pass






   def     _SCREAMING_SNAKE_CASE    (     self			:					str  ):

     '''simple docstring'''
     __A      ,   __A						      =			self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       with self.subTest(model_class.__name__  ):
         __A						      =			self._prepare_for_class(_lowerCamelCase,	_lowerCamelCase  )
         __A						      =			model_class(_lowerCamelCase  )

         @jax.jit
         def model_jitted(_lowerCamelCase			:					Optional[int],	**_lowerCamelCase			:					Tuple  ):
           return model(pixel_values=_lowerCamelCase,	**_lowerCamelCase  )

         with self.subTest('''JIT Enabled'''  ):
           __A						      =			model_jitted(**_lowerCamelCase  ).to_tuple()

         with self.subTest('''JIT Disabled'''  ):
           with jax.disable_jit():
             __A						      =			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    (   ):




  """simple docstring"""







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





@require_flax
class  snake_case		(							unittest.TestCase						):


   '''simple docstring'''

   @cached_property
   def     _SCREAMING_SNAKE_CASE    (     self			:					List[str]  ):

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






   @slow
   def     _SCREAMING_SNAKE_CASE    (     self			:					Tuple  ):

     '''simple docstring'''
     __A						      =			FlaxResNetForImageClassification.from_pretrained('''microsoft/resnet-50'''  )

     __A						      =			self.default_image_processor
     __A						      =			prepare_img()
     __A						      =			image_processor(images=_lowerCamelCase,	return_tensors='''np'''  )

     __A						      =			model(**_lowerCamelCase  )

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

     __A						      =			jnp.array([-11.10_69, -9.78_77, -8.37_77]  )

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












"""simple docstring"""


# This test is meant to be run in on an instance with TPUs like this:
#
#   python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
#
# Replace 8 with the number of TPU cores you have.
#

import sys
from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.utils import logging


lowercase_    					=					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  snake_case		(							_lowerCAmelCase						):


          '''simple docstring'''

          def __init__(     self			:					List[Any],	_lowerCamelCase			:					int = 1_01  ):

            '''simple docstring'''
            __A						      =			length


          def __len__(     self			:					Optional[int]  ):

            '''simple docstring'''
            return self.length






          def __getitem__(     self			:					List[str],	_lowerCamelCase			:					Union[str, Any]  ):

            '''simple docstring'''
            return i





       class  snake_case		:


          '''simple docstring'''

          def __call__(     self			:					Tuple,	_lowerCamelCase			:					Dict  ):

            '''simple docstring'''
            return {"input_ids": torch.tensor(_lowerCamelCase  ), "labels": torch.tensor(_lowerCamelCase  )}





       class  snake_case		(							nn.Module						):


          '''simple docstring'''

          def __init__(     self			:					List[str]  ):

            '''simple docstring'''
            super().__init__()
            # Add some (unused) params otherwise DDP will complain.
            __A						      =			nn.Linear(1_20,	80  )






          def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[Any],	_lowerCamelCase			:					Tuple,	_lowerCamelCase			:					Optional[Any]=None  ):

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



def  lowerCAmelCase    (   ):




  """simple docstring"""







  __A						      =			HfArgumentParser((TrainingArguments,)			)
  sys.argv += ["--output_dir", "./examples"]
  __A						      =			parser.parse_args_into_dataclasses()[0]

  logger.warning(
      f'Process rank: {training_args.local_rank}, device: {training_args.device}, '
      f'tpu_num_cores: {training_args.tpu_num_cores}'		,	)

  # 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 [1_0_0_1, 2_5_6, 1_5]:
    __A						      =			DummyDataset(__UpperCamelCase			)

    def compute_metrics(__UpperCamelCase			) -> Dict:
      __A						      =			list(range(len(__UpperCamelCase			)			)			)
      __A						      =			p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
      return {"success": success}

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

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

    __A						      =			2

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

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

    __A						      =			None

  logger.info('''🔥 All distributed tests successful'''			)



def  lowerCAmelCase    (   __UpperCamelCase			):




  """simple docstring"""







  main()


if __name__ == "__main__":
       main()




