





from __future__ import annotations

import inspect
import unittest
from math import floor

import numpy as np

from transformers import CvtConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import TFCvtForImageClassification, TFCvtModel
    from transformers.models.cvt.modeling_tf_cvt import TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor



class       A__     ( __snake_case    ):





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Optional[Any] 			=     self.config_class(**self.inputs_dict     )
       self.parent.assertTrue(hasattr(A_ ,							"embed_dim"     )     )
       self.parent.assertTrue(hasattr(A_ ,							"num_heads"     )     )



class       A__     :





      def __init__(				self ,							A_ ,							A_=13 ,							A_=64 ,							A_=3 ,							A_=[16, 48, 96] ,							A_=[1, 3, 6] ,							A_=[1, 2, 10] ,							A_=[7, 3, 3] ,							A_=[4, 2, 2] ,							A_=[2, 1, 1] ,							A_=[2, 2, 2] ,							A_=[False, False, True] ,							A_=[0.0, 0.0, 0.0] ,							A_=0.02 ,							A_=1e-12 ,							A_=True ,							A_=True ,							A_=2 ,							):
       '''simple docstring'''




       UpperCamelCase			:      int 			=     parent
       UpperCamelCase			:      int 			=     batch_size
       UpperCamelCase			:      Dict 			=     image_size
       UpperCamelCase			:      Tuple 			=     patch_sizes
       UpperCamelCase			:      Optional[Any] 			=     patch_stride
       UpperCamelCase			:      Tuple 			=     patch_padding
       UpperCamelCase			:      Union[str, Any] 			=     is_training
       UpperCamelCase			:      int 			=     use_labels
       UpperCamelCase			:      List[Any] 			=     num_labels
       UpperCamelCase			:      Any 			=     num_channels
       UpperCamelCase			:      Optional[Any] 			=     embed_dim
       UpperCamelCase			:      int 			=     num_heads
       UpperCamelCase			:      Union[str, Any] 			=     stride_kv
       UpperCamelCase			:      Optional[int] 			=     depth
       UpperCamelCase			:      Any 			=     cls_token
       UpperCamelCase			:      List[Any] 			=     attention_drop_rate
       UpperCamelCase			:      Optional[Any] 			=     initializer_range
       UpperCamelCase			:      Any 			=     layer_norm_eps





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




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

       UpperCamelCase			:      Optional[int] 			=     None
       if self.use_labels:
        # create a random int32 tensor of given shape
        UpperCamelCase			:      Tuple 			=     ids_tensor([self.batch_size] ,							self.num_labels     )

       UpperCamelCase			:      str 			=     self.get_config()
       return config, pixel_values, labels





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       return CvtConfig(
           image_size=self.image_size ,							num_labels=self.num_labels ,							num_channels=self.num_channels ,							embed_dim=self.embed_dim ,							num_heads=self.num_heads ,							patch_sizes=self.patch_sizes ,							patch_padding=self.patch_padding ,							patch_stride=self.patch_stride ,							stride_kv=self.stride_kv ,							depth=self.depth ,							cls_token=self.cls_token ,							attention_drop_rate=self.attention_drop_rate ,							initializer_range=self.initializer_range ,							)





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      Any 			=     TFCvtModel(config=A_     )
       UpperCamelCase			:      Any 			=     model(A_ ,							training=A_     )
       UpperCamelCase			:      Union[str, Any] 			=     (self.image_size, self.image_size)
       UpperCamelCase			,					UpperCamelCase			:      List[str] 			=     image_size[0], image_size[1]
       for i in range(len(self.depth     )     ):
        UpperCamelCase			:      str 			=     floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1     )
        UpperCamelCase			:      Any 			=     floor(((width + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1     )
       self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, self.embed_dim[-1], height, width)     )





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      Optional[Any] 			=     self.num_labels
       UpperCamelCase			:      Optional[Any] 			=     TFCvtForImageClassification(A_     )
       UpperCamelCase			:      Optional[int] 			=     model(A_ ,							labels=A_ ,							training=A_     )
       self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.num_labels)     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




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



@require_tf
class       A__     ( __snake_case ,     __snake_case ,     unittest.TestCase    ):
      _UpperCAmelCase			:int							=     (TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
      _UpperCAmelCase			:Any							=     (
          {'feature-extraction': TFCvtModel, 'image-classification': TFCvtForImageClassification}
          if is_tf_available()
          else {}
      )
      _UpperCAmelCase			:Optional[int]							=     False
      _UpperCAmelCase			:str							=     False
      _UpperCAmelCase			:Union[str, Any]							=     False
      _UpperCAmelCase			:str							=     False
      _UpperCAmelCase			:Tuple							=     False





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Union[str, Any] 			=     TFCvtModelTester(self     )
       UpperCamelCase			:      Union[str, Any] 			=     TFCvtConfigTester(self ,							config_class=A_ ,							has_text_modality=A_ ,							hidden_size=37     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       self.config_tester.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()





      @unittest.skip(reason="Cvt does not output attentions"     )
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       pass





      @unittest.skip(reason="Cvt does not use inputs_embeds"     )
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       pass





      @unittest.skip(reason="Cvt does not support input and output embeddings"     )
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       pass





      @unittest.skipIf(
          not is_tf_available() or len(tf.config.list_physical_devices("GPU"     )     ) == 0 ,							reason="TF does not support backprop for grouped convolutions on CPU." ,							)
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       super().test_dataset_conversion()





      @unittest.skipIf(
          not is_tf_available() or len(tf.config.list_physical_devices("GPU"     )     ) == 0 ,							reason="TF does not support backprop for grouped convolutions on CPU." ,							)
      @slow
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       super().test_keras_fit()





      @unittest.skip(reason="Get `Failed to determine best cudnn convolution algo.` error after using TF 2.12+cuda 11.8"     )
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      int 			=     tf.keras.mixed_precision.Policy("mixed_float16"     )
       tf.keras.mixed_precision.set_global_policy(A_     )
       super().test_keras_fit()
       tf.keras.mixed_precision.set_global_policy("float32"     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




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

       for model_class in self.all_model_classes:
        UpperCamelCase			:      Optional[Any] 			=     model_class(A_     )
        UpperCamelCase			:      List[Any] 			=     inspect.signature(model.call     )
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        UpperCamelCase			:      int 			=     [*signature.parameters.keys()]

        UpperCamelCase			:      Optional[int] 			=     ["pixel_values"]
        self.assertListEqual(arg_names[:1] ,							A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       def check_hidden_states_output(A_ ,							A_ ,							A_     ):
        UpperCamelCase			:      str 			=     model_class(A_     )

        UpperCamelCase			:      List[str] 			=     model(**self._prepare_for_class(A_ ,							A_     )     )
        UpperCamelCase			:      Dict 			=     outputs.hidden_states

        UpperCamelCase			:      int 			=     len(self.model_tester.depth     )
        self.assertEqual(len(A_     ) ,							A_     )

        # verify the first hidden states (first block)
        self.assertListEqual(
            list(hidden_states[0].shape[-3:]     ) ,							[
                self.model_tester.embed_dim[0],
                self.model_tester.image_size // 4,
                self.model_tester.image_size // 4,
            ] ,							)

       UpperCamelCase			,					UpperCamelCase			:      str 			=     self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCamelCase			:      Union[str, Any] 			=     True
        check_hidden_states_output(A_ ,							A_ ,							A_     )

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        UpperCamelCase			:      Dict 			=     True

        check_hidden_states_output(A_ ,							A_ ,							A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Tuple 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*A_     )





      @slow
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCamelCase			:      List[str] 			=     TFCvtModel.from_pretrained(A_     )
        self.assertIsNotNone(A_     )



def 	A_       (  )		->				Dict:
 UpperCamelCase			:      List[Any] 			=     Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"  )
 return image



@require_tf
@require_vision
class       A__     ( unittest.TestCase    ):





      @cached_property
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]     )





      @slow
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]     )

       UpperCamelCase			:      Optional[int] 			=     self.default_image_processor
       UpperCamelCase			:      List[str] 			=     prepare_img()
       UpperCamelCase			:      Union[str, Any] 			=     image_processor(images=A_ ,							return_tensors="tf"     )

       # forward pass
       UpperCamelCase			:      Dict 			=     model(**A_     )

       # verify the logits
       UpperCamelCase			:      Optional[Any] 			=     tf.TensorShape((1, 1000)     )
       self.assertEqual(outputs.logits.shape ,							A_     )

       UpperCamelCase			:      int 			=     tf.constant([0.92_85, 0.90_15, -0.31_50]     )
       self.assertTrue(np.allclose(outputs.logits[0, :3].numpy() ,							A_ ,							atol=1e-4     )     )







import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp



def 	A_       (  )		->				Dict:
 UpperCamelCase			:      Tuple 			=     ArgumentParser(
     description=(
         "PyTorch TPU distributed training launch "
         "helper utility that will spawn up "
         "multiple distributed processes"
     )  )

 # Optional arguments for the launch helper
 parser.add_argument("--num_cores"							,     type=_lowerCAmelCase							,     default=1							,     help="Number of TPU cores to use (1 or 8)."  )

 # positional
 parser.add_argument(
     "training_script"							,     type=_lowerCAmelCase							,     help=(
         "The full path to the single TPU training "
         "program/script to be launched in parallel, "
         "followed by all the arguments for the "
         "training script"
     )							,     )

 # rest from the training program
 parser.add_argument("training_script_args"							,     nargs=_lowerCAmelCase  )

 return parser.parse_args()



def 	A_       (  )		->				Optional[int]:
 UpperCamelCase			:      Tuple 			=     parse_args()

 # Import training_script as a module.
 UpperCamelCase			:      Union[str, Any] 			=     Path(args.training_script  )
 sys.path.append(str(script_fpath.parent.resolve()  )  )
 UpperCamelCase			:      List[Any] 			=     script_fpath.stem
 UpperCamelCase			:      Optional[Any] 			=     importlib.import_module(_lowerCAmelCase  )

 # Patch sys.argv
 UpperCamelCase			:      List[Any] 			=     [args.training_script] + args.training_script_args + ["--tpu_num_cores", str(args.num_cores  )]
 xmp.spawn(mod._mp_fn							,     args=()							,     nprocs=args.num_cores  )


if __name__ == "__main__":
    main()
