


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  SCREAMING_SNAKE_CASE__  (				lowercase__			):
      def        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     )		->			List[str]:
             a_					:		Any							=  self.config_class(**self.inputs_dict     )
             self.parent.assertTrue(hasattr(SCREAMING_SNAKE_CASE__     ,							'embed_dim'     )     )
             self.parent.assertTrue(hasattr(SCREAMING_SNAKE_CASE__     ,							'num_heads'     )     )
class  SCREAMING_SNAKE_CASE__  :
      def __init__(       self				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       Tuple=1_3     ,							SCREAMING_SNAKE_CASE__				:       Tuple=6_4     ,							SCREAMING_SNAKE_CASE__				:       Dict=3     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]=[1_6, 4_8, 9_6]     ,							SCREAMING_SNAKE_CASE__				:       str=[1, 3, 6]     ,							SCREAMING_SNAKE_CASE__				:       Tuple=[1, 2, 1_0]     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]=[7, 3, 3]     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]=[4, 2, 2]     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]=[2, 1, 1]     ,							SCREAMING_SNAKE_CASE__				:       Dict=[2, 2, 2]     ,							SCREAMING_SNAKE_CASE__				:       str=[False, False, True]     ,							SCREAMING_SNAKE_CASE__				:       Tuple=[0.0, 0.0, 0.0]     ,							SCREAMING_SNAKE_CASE__				:       Dict=0.02     ,							SCREAMING_SNAKE_CASE__				:       Dict=1E-12     ,							SCREAMING_SNAKE_CASE__				:       List[str]=True     ,							SCREAMING_SNAKE_CASE__				:       Tuple=True     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]=2     ,							)		->			int:
             a_					:		Optional[Any]							=  parent
             a_					:		List[str]							=  batch_size
             a_					:		Optional[Any]							=  image_size
             a_					:		Tuple							=  patch_sizes
             a_					:		int							=  patch_stride
             a_					:		Union[str, Any]							=  patch_padding
             a_					:		List[Any]							=  is_training
             a_					:		Optional[Any]							=  use_labels
             a_					:		Tuple							=  num_labels
             a_					:		int							=  num_channels
             a_					:		Dict							=  embed_dim
             a_					:		int							=  num_heads
             a_					:		Tuple							=  stride_kv
             a_					:		Any							=  depth
             a_					:		List[str]							=  cls_token
             a_					:		Any							=  attention_drop_rate
             a_					:		List[Any]							=  initializer_range
             a_					:		List[Any]							=  layer_norm_eps
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			str:
             a_					:		Dict							=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

             a_					:		Union[str, Any]							=  None
             if self.use_labels:
                    # create a random int32 tensor of given shape
                    a_					:		int							=  ids_tensor([self.batch_size]     ,							self.num_labels     )

             a_					:		int							=  self.get_config()
             return config, pixel_values, labels
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			Optional[Any]:
             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        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]     )		->			Optional[int]:
             a_					:		Dict							=  TFCvtModel(config=SCREAMING_SNAKE_CASE__     )
             a_					:		Union[str, Any]							=  model(SCREAMING_SNAKE_CASE__     ,							training=SCREAMING_SNAKE_CASE__     )
             a_					:		List[Any]							=  (self.image_size, self.image_size)
             a_							,		a_					:		Union[str, Any]							=  image_size[0], image_size[1]
             for i in range(len(self.depth     )     ):
                    a_					:		Any							=  floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1     )
                    a_					:		List[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        SCREAMING_SNAKE_CASE							(       self				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       Any     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       Any     )		->			Optional[Any]:
             a_					:		List[str]							=  self.num_labels
             a_					:		List[str]							=  TFCvtForImageClassification(SCREAMING_SNAKE_CASE__     )
             a_					:		int							=  model(SCREAMING_SNAKE_CASE__     ,							labels=SCREAMING_SNAKE_CASE__     ,							training=SCREAMING_SNAKE_CASE__     )
             self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.num_labels)     )

      def        SCREAMING_SNAKE_CASE							(       self				:       Union[str, Any]     )		->			Optional[int]:
             a_					:		Optional[Any]							=  self.prepare_config_and_inputs()
             a_							,		a_							,		a_					:		Union[str, Any]							=  config_and_inputs
             a_					:		Any							=  {'pixel_values': pixel_values}
             return config, inputs_dict
@require_tf
class  SCREAMING_SNAKE_CASE__  (				lowercase__						,    lowercase__						,    unittest.TestCase			):
      snake_case__				:     str        =						(TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
      snake_case__				:     List[Any]        =						(
          {'''feature-extraction''': TFCvtModel, '''image-classification''': TFCvtForImageClassification}
          if is_tf_available()
          else {}
      )
      snake_case__				:     List[Any]        =						False
      snake_case__				:     Dict        =						False
      snake_case__				:     Optional[Any]        =						False
      snake_case__				:     Optional[int]        =						False
      snake_case__				:     List[Any]        =						False
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			int:
             a_					:		str							=  TFCvtModelTester(self     )
             a_					:		List[Any]							=  TFCvtConfigTester(self     ,							config_class=SCREAMING_SNAKE_CASE__     ,							has_text_modality=SCREAMING_SNAKE_CASE__     ,							hidden_size=3_7     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[int]     )		->			List[Any]:
             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        SCREAMING_SNAKE_CASE							(       self				:       int     )		->			str:
             pass
      @unittest.skip(reason='Cvt does not use inputs_embeds'     )
      def        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     )		->			Optional[int]:
             pass
      @unittest.skip(reason='Cvt does not support input and output embeddings'     )
      def        SCREAMING_SNAKE_CASE							(       self				:       int     )		->			Union[str, Any]:
             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        SCREAMING_SNAKE_CASE							(       self				:       int     )		->			int:
             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        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Any:
             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        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			Any:
             a_					:		Optional[int]							=  tf.keras.mixed_precision.Policy('mixed_float16'     )
             tf.keras.mixed_precision.set_global_policy(SCREAMING_SNAKE_CASE__     )
             super().test_keras_fit()
             tf.keras.mixed_precision.set_global_policy('float32'     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			List[Any]:
             a_							,		a_					:		int							=  self.model_tester.prepare_config_and_inputs_for_common()

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

                    a_					:		Union[str, Any]							=  ['pixel_values']
                    self.assertListEqual(arg_names[:1]     ,							SCREAMING_SNAKE_CASE__     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Union[str, Any]     )		->			Optional[Any]:
             def check_hidden_states_output(SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ):
                    a_					:		Union[str, Any]							=  model_class(SCREAMING_SNAKE_CASE__     )

                    a_					:		Optional[int]							=  model(**self._prepare_for_class(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     )     )
                    a_					:		Optional[int]							=  outputs.hidden_states

                    a_					:		Dict							=  len(self.model_tester.depth     )
                    self.assertEqual(len(SCREAMING_SNAKE_CASE__     )     ,							SCREAMING_SNAKE_CASE__     )

                    # 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,
                        ]     ,							)

             a_							,		a_					:		Any							=  self.model_tester.prepare_config_and_inputs_for_common()

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

                    # check that output_hidden_states also work using config
                    del inputs_dict["output_hidden_states"]
                    a_					:		Union[str, Any]							=  True

                    check_hidden_states_output(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Union[str, Any]     )		->			int:
             a_					:		str							=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE__     )
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			str:
             a_					:		int							=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE__     )

      @slow
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[int]     )		->			Any:
             for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    a_					:		Any							=  TFCvtModel.from_pretrained(SCREAMING_SNAKE_CASE__     )
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE__     )







def        SCREAMING_SNAKE_CASE_						(			)	->		Optional[int]:

       """simple docstring"""

       a_					:		Optional[int]							=  Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
       return image
@require_tf
@require_vision
class  SCREAMING_SNAKE_CASE__  (				unittest.TestCase			):
      @cached_property
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Dict:
             return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]     )

      @slow
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			Optional[int]:
             a_					:		int							=  TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]     )

             a_					:		Optional[int]							=  self.default_image_processor
             a_					:		List[Any]							=  prepare_img()
             a_					:		Optional[Any]							=  image_processor(images=SCREAMING_SNAKE_CASE__     ,							return_tensors='tf'     )

             # forward pass
             a_					:		Optional[Any]							=  model(**SCREAMING_SNAKE_CASE__     )

             # verify the logits
             a_					:		List[Any]							=  tf.TensorShape((1, 1_0_0_0)     )
             self.assertEqual(outputs.logits.shape     ,							SCREAMING_SNAKE_CASE__     )

             a_					:		Any							=  tf.constant([0.9285, 0.9015, -0.3150]     )
             self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()     ,							SCREAMING_SNAKE_CASE__     ,							atol=1E-4     )     )






import flax.linen as nn
import jax
import jax.numpy as jnp
class  SCREAMING_SNAKE_CASE__  (				nn.Module			):
      snake_case__				:     int
      snake_case__				:     jnp.dtype        =						jnp.floataa
      def        SCREAMING_SNAKE_CASE							(       self				:       str     )		->			int:
             a_					:		Dict							=  nn.Conv(
                 self.out_channels     ,							kernel_size=(3, 3)     ,							strides=(1, 1)     ,							padding=((1, 1), (1, 1))     ,							dtype=self.dtype     ,							)

      def __call__(       self				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     )		->			Optional[int]:
             a_							,		a_							,		a_							,		a_					:		Union[str, Any]							=  hidden_states.shape
             a_					:		List[str]							=  jax.image.resize(
                 SCREAMING_SNAKE_CASE__     ,							shape=(batch, height * 2, width * 2, channels)     ,							method='nearest'     ,							)
             a_					:		Any							=  self.conv(SCREAMING_SNAKE_CASE__     )
             return hidden_states
class  SCREAMING_SNAKE_CASE__  (				nn.Module			):
      snake_case__				:     int
      snake_case__				:     jnp.dtype        =						jnp.floataa
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Optional[int]:
             a_					:		Optional[int]							=  nn.Conv(
                 self.out_channels     ,							kernel_size=(3, 3)     ,							strides=(2, 2)     ,							padding=((1, 1), (1, 1))     ,							dtype=self.dtype     ,							)

      def __call__(       self				:       Any     ,							SCREAMING_SNAKE_CASE__				:       int     )		->			Optional[int]:
             # pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
             # hidden_states = jnp.pad(hidden_states, pad_width=pad)
             a_					:		str							=  self.conv(SCREAMING_SNAKE_CASE__     )
             return hidden_states
class  SCREAMING_SNAKE_CASE__  (				nn.Module			):
      snake_case__				:     int
      snake_case__				:     int        =						None
      snake_case__				:     float        =						0.0
      snake_case__				:     bool        =						None
      snake_case__				:     jnp.dtype        =						jnp.floataa
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Dict:
             a_					:		List[str]							=  self.in_channels if self.out_channels is None else self.out_channels

             a_					:		Optional[int]							=  nn.GroupNorm(num_groups=3_2     ,							epsilon=1E-5     )
             a_					:		Any							=  nn.Conv(
                 SCREAMING_SNAKE_CASE__     ,							kernel_size=(3, 3)     ,							strides=(1, 1)     ,							padding=((1, 1), (1, 1))     ,							dtype=self.dtype     ,							)

             a_					:		Optional[int]							=  nn.Dense(SCREAMING_SNAKE_CASE__     ,							dtype=self.dtype     )

             a_					:		Union[str, Any]							=  nn.GroupNorm(num_groups=3_2     ,							epsilon=1E-5     )
             a_					:		int							=  nn.Dropout(self.dropout_prob     )
             a_					:		Optional[Any]							=  nn.Conv(
                 SCREAMING_SNAKE_CASE__     ,							kernel_size=(3, 3)     ,							strides=(1, 1)     ,							padding=((1, 1), (1, 1))     ,							dtype=self.dtype     ,							)

             a_					:		List[str]							=  self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

             a_					:		List[Any]							=  None
             if use_nin_shortcut:
                    a_					:		Union[str, Any]							=  nn.Conv(
                        SCREAMING_SNAKE_CASE__     ,							kernel_size=(1, 1)     ,							strides=(1, 1)     ,							padding='VALID'     ,							dtype=self.dtype     ,							)

      def __call__(       self				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       List[Any]=True     )		->			int:
             a_					:		List[Any]							=  hidden_states
             a_					:		Any							=  self.norma(SCREAMING_SNAKE_CASE__     )
             a_					:		Any							=  nn.swish(SCREAMING_SNAKE_CASE__     )
             a_					:		Union[str, Any]							=  self.conva(SCREAMING_SNAKE_CASE__     )

             a_					:		int							=  self.time_emb_proj(nn.swish(SCREAMING_SNAKE_CASE__     )     )
             a_					:		List[str]							=  jnp.expand_dims(jnp.expand_dims(SCREAMING_SNAKE_CASE__     ,							1     )     ,							1     )
             a_					:		Optional[int]							=  hidden_states + temb

             a_					:		List[str]							=  self.norma(SCREAMING_SNAKE_CASE__     )
             a_					:		Tuple							=  nn.swish(SCREAMING_SNAKE_CASE__     )
             a_					:		Optional[int]							=  self.dropout(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     )
             a_					:		List[Any]							=  self.conva(SCREAMING_SNAKE_CASE__     )

             if self.conv_shortcut is not None:
                    a_					:		List[str]							=  self.conv_shortcut(SCREAMING_SNAKE_CASE__     )

             return hidden_states + residual


