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


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





      """simple docstring"""





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


   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      _lowercase    :   str      =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

      _lowercase    :   Tuple      =     self.get_config()

      return config, pixel_values


   def 		_lowerCamelCase     (				self						):





      """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    ,       _UpperCamelCase    ,       _UpperCamelCase						):





      """simple docstring"""





      _lowercase    :   Tuple      =     FlaxResNetModel(config=_UpperCamelCase						)
      _lowercase    :   Any      =     model(_UpperCamelCase						)

      # 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    ,       _UpperCamelCase    ,       _UpperCamelCase						):





      """simple docstring"""





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






   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





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




@require_flax
class      a__						(							lowerCamelCase_						,   unittest.TestCase						):
   _SCREAMING_SNAKE_CASE					:  str			 =							(FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

   _SCREAMING_SNAKE_CASE					:  Tuple			 =							False
   _SCREAMING_SNAKE_CASE					:  Optional[Any]			 =							False
   _SCREAMING_SNAKE_CASE					:  Dict			 =							False


   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      _lowercase    :   str      =     FlaxResNetModelTester(self						)
      _lowercase    :   List[str]      =     ConfigTester(self    ,       config_class=_UpperCamelCase    ,       has_text_modality=_UpperCamelCase						)


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





      """simple docstring"""





      return


   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      _lowercase    :   List[Any]      =     self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*_UpperCamelCase						)


   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      _lowercase    :   List[str]      =     self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*_UpperCamelCase						)


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





      """simple docstring"""





      pass


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





      """simple docstring"""





      pass


   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      _lowercase    :   List[Any]      =     self.model_tester.prepare_config_and_inputs_for_common()

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

         _lowercase    :   str      =     ["pixel_values"]
         self.assertListEqual(arg_names[:1]    ,       _UpperCamelCase						)


   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      def check_hidden_states_output(_UpperCamelCase    ,       _UpperCamelCase    ,       _UpperCamelCase						):
         _lowercase    :   List[Any]      =     model_class(_UpperCamelCase						)

         _lowercase    :   Optional[Any]      =     model(**self._prepare_for_class(_UpperCamelCase    ,       _UpperCamelCase						)						)

         _lowercase    :   Union[str, Any]      =     outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

         _lowercase    :   int      =     self.model_tester.num_stages
         self.assertEqual(len(_UpperCamelCase						)    ,       expected_num_stages + 1						)


   @unittest.skip(reason="ResNet does not use feedforward chunking"						)
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      pass






   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      _lowercase    :   List[Any]      =     self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
         with self.subTest(model_class.__name__						):
            _lowercase    :   Dict      =     self._prepare_for_class(_UpperCamelCase    ,       _UpperCamelCase						)
            _lowercase    :   Dict      =     model_class(_UpperCamelCase						)

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

            with self.subTest("JIT Enabled"						):
               _lowercase    :   List[str]      =     model_jitted(**_UpperCamelCase						).to_tuple()

            with self.subTest("JIT Disabled"						):
               with jax.disable_jit():
                  _lowercase    :   str      =     model_jitted(**_UpperCamelCase						).to_tuple()

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







def       _A	(			)    ->			int:
   _lowercase    :   Tuple      =     Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"  )
   return image


@require_flax
class      a__						(							unittest.TestCase						):


   @cached_property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





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






   @slow
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      _lowercase    :   Union[str, Any]      =     FlaxResNetForImageClassification.from_pretrained("microsoft/resnet-50"						)

      _lowercase    :   List[Any]      =     self.default_image_processor
      _lowercase    :   List[Any]      =     prepare_img()
      _lowercase    :   Optional[int]      =     image_processor(images=_UpperCamelCase    ,       return_tensors="np"						)

      _lowercase    :   Union[str, Any]      =     model(**_UpperCamelCase						)

      # verify the logits
      _lowercase    :   Optional[Any]      =     (1, 1000)
      self.assertEqual(outputs.logits.shape    ,       _UpperCamelCase						)

      _lowercase    :   Union[str, Any]      =     jnp.array([-11.1069, -9.7_8_7_7, -8.3_7_7_7]						)

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









'''simple docstring'''
import warnings
from dataclasses import dataclass, field
from typing import Optional, Tuple

from .training_args import TrainingArguments
from .utils import cached_property, is_tf_available, logging, requires_backends


_snake_case     						=						logging.get_logger(__name__)

if is_tf_available():
    import tensorflow as tf


@dataclass
class      a__						(							lowerCamelCase_						):
   _SCREAMING_SNAKE_CASE					:  int			 =							'tf'
   _SCREAMING_SNAKE_CASE					:  Optional[str]			 =							field(
       default=lowerCamelCase_						,   metadata={'help': 'Name of TPU'}						,   )

   _SCREAMING_SNAKE_CASE					:  Optional[str]			 =							field(
       default=lowerCamelCase_						,   metadata={'help': 'Zone of TPU'}						,   )

   _SCREAMING_SNAKE_CASE					:  Optional[str]			 =							field(
       default=lowerCamelCase_						,   metadata={'help': 'Name of Cloud TPU-enabled project'}						,   )

   _SCREAMING_SNAKE_CASE					:  float			 =							field(
       default=1.0						,   metadata={'help': 'Power for the Polynomial decay LR scheduler.'}						,   )

   _SCREAMING_SNAKE_CASE					:  bool			 =							field(default=lowerCamelCase_						,   metadata={'help': 'Whether to activate the XLA compilation or not'}						)


   @cached_property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      requires_backends(self    ,       ["tf"]						)
      logger.info("Tensorflow: setting up strategy"						)

      _lowercase    :   Optional[int]      =     tf.config.list_physical_devices("GPU"						)

      # Set to float16 at first
      if self.fpaa:
         tf.keras.mixed_precision.set_global_policy("mixed_float16"						)

      if self.no_cuda:
         _lowercase    :   Optional[int]      =     tf.distribute.OneDeviceStrategy(device="/cpu:0"						)
      else:
         try:
            if self.tpu_name:
               _lowercase    :   int      =     tf.distribute.cluster_resolver.TPUClusterResolver(
                   self.tpu_name    ,       zone=self.tpu_zone    ,       project=self.gcp_project						)
            else:
               _lowercase    :   Any      =     tf.distribute.cluster_resolver.TPUClusterResolver()
         except ValueError:
            if self.tpu_name:
               raise RuntimeError(f'''Couldn\'t connect to TPU {self.tpu_name}!'''						)
            else:
               _lowercase    :   Tuple      =     None

         if tpu:
            # Set to bfloat16 in case of TPU
            if self.fpaa:
               tf.keras.mixed_precision.set_global_policy("mixed_bfloat16"						)

            tf.config.experimental_connect_to_cluster(_UpperCamelCase						)
            tf.tpu.experimental.initialize_tpu_system(_UpperCamelCase						)

            _lowercase    :   List[str]      =     tf.distribute.TPUStrategy(_UpperCamelCase						)

         elif len(_UpperCamelCase						) == 0:
            _lowercase    :   Optional[Any]      =     tf.distribute.OneDeviceStrategy(device="/cpu:0"						)
         elif len(_UpperCamelCase						) == 1:
            _lowercase    :   Optional[int]      =     tf.distribute.OneDeviceStrategy(device="/gpu:0"						)
         elif len(_UpperCamelCase						) > 1:
            # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
            _lowercase    :   Any      =     tf.distribute.MirroredStrategy()
         else:
            raise ValueError("Cannot find the proper strategy, please check your environment properties."						)

      return strategy


   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      requires_backends(self    ,       ["tf"]						)
      return self._setup_strategy


   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      requires_backends(self    ,       ["tf"]						)
      return self._setup_strategy.num_replicas_in_sync


   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      return False  # TF Logging is handled by Keras not the Trainer


   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      if self.per_gpu_train_batch_size:
         logger.warning(
             "Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future "
             "version. Using `--per_device_train_batch_size` is preferred."						)
      _lowercase    :   str      =     self.per_gpu_train_batch_size or self.per_device_train_batch_size
      return per_device_batch_size * self.n_replicas


   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      if self.per_gpu_eval_batch_size:
         logger.warning(
             "Using deprecated `--per_gpu_eval_batch_size` argument which will be removed in a future "
             "version. Using `--per_device_eval_batch_size` is preferred."						)
      _lowercase    :   List[Any]      =     self.per_gpu_eval_batch_size or self.per_device_eval_batch_size
      return per_device_batch_size * self.n_replicas






   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      requires_backends(self    ,       ["tf"]						)
      warnings.warn(
          "The n_gpu argument is deprecated and will be removed in a future version, use n_replicas instead."    ,       _UpperCamelCase    ,       )
      return self._setup_strategy.num_replicas_in_sync






