


'''simple docstring'''






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    __snake_case       (   __a):







    """simple docstring"""



    def __lowercase   (   self    :  Union[str, Any]		)   ->  Any:
      lowerCAmelCase_			:      str			=   self.config_class(**self.inputs_dict		)
      self.parent.assertTrue(hasattr(_A    ,						"""embed_dim"""		)		)
      self.parent.assertTrue(hasattr(_A    ,						"""num_heads"""		)		)







class    __snake_case       :







    """simple docstring"""



    def __init__(   self    :  int    ,						lowerCamelCase    :  Tuple    ,						lowerCamelCase    :  Any=13    ,						lowerCamelCase    :  Optional[int]=64    ,						lowerCamelCase    :  Optional[Any]=3    ,						lowerCamelCase    :  List[str]=[16, 48, 96]    ,						lowerCamelCase    :  int=[1, 3, 6]    ,						lowerCamelCase    :  Optional[int]=[1, 2, 10]    ,						lowerCamelCase    :  int=[7, 3, 3]    ,						lowerCamelCase    :  Union[str, Any]=[4, 2, 2]    ,						lowerCamelCase    :  Dict=[2, 1, 1]    ,						lowerCamelCase    :  Optional[Any]=[2, 2, 2]    ,						lowerCamelCase    :  Optional[Any]=[False, False, True]    ,						lowerCamelCase    :  List[Any]=[0.0, 0.0, 0.0]    ,						lowerCamelCase    :  str=0.02    ,						lowerCamelCase    :  Tuple=1E-12    ,						lowerCamelCase    :  Union[str, Any]=True    ,						lowerCamelCase    :  Optional[Any]=True    ,						lowerCamelCase    :  Optional[int]=2    ,						)   ->  Union[str, Any]:
      lowerCAmelCase_			:      Dict			=   parent
      lowerCAmelCase_			:      List[str]			=   batch_size
      lowerCAmelCase_			:      Optional[int]			=   image_size
      lowerCAmelCase_			:      List[str]			=   patch_sizes
      lowerCAmelCase_			:      Any			=   patch_stride
      lowerCAmelCase_			:      Tuple			=   patch_padding
      lowerCAmelCase_			:      int			=   is_training
      lowerCAmelCase_			:      Dict			=   use_labels
      lowerCAmelCase_			:      List[Any]			=   num_labels
      lowerCAmelCase_			:      Optional[Any]			=   num_channels
      lowerCAmelCase_			:      Optional[int]			=   embed_dim
      lowerCAmelCase_			:      int			=   num_heads
      lowerCAmelCase_			:      Any			=   stride_kv
      lowerCAmelCase_			:      str			=   depth
      lowerCAmelCase_			:      List[Any]			=   cls_token
      lowerCAmelCase_			:      List[Any]			=   attention_drop_rate
      lowerCAmelCase_			:      Optional[int]			=   initializer_range
      lowerCAmelCase_			:      Optional[int]			=   layer_norm_eps

    def __lowercase   (   self    :  Optional[int]		)   ->  str:
      lowerCAmelCase_			:      int			=   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

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

      lowerCAmelCase_			:      List[Any]			=   self.get_config()
      return config, pixel_values, labels

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





    def __lowercase   (   self    :  Optional[int]		)   ->  Dict:
      lowerCAmelCase_			:      List[Any]			=   self.prepare_config_and_inputs()
      lowerCAmelCase_			:      Dict			=   config_and_inputs
      lowerCAmelCase_			:      Optional[Any]			=   {'''pixel_values''': pixel_values}
      return config, inputs_dict







@require_tf
class    __snake_case       (   __a	,__a	,unittest.TestCase):







    """simple docstring"""



    lowercase    =					(TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
    lowercase    =					(
        {'feature-extraction': TFCvtModel, 'image-classification': TFCvtForImageClassification}
        if is_tf_available()
        else {}
    )
    lowercase    =					False
    lowercase    =					False
    lowercase    =					False
    lowercase    =					False
    lowercase    =					False

    def __lowercase   (   self    :  Any		)   ->  str:
      lowerCAmelCase_			:      Union[str, Any]			=   TFCvtModelTester(self		)
      lowerCAmelCase_			:      Tuple			=   TFCvtConfigTester(self    ,						config_class=_A    ,						has_text_modality=_A    ,						hidden_size=37		)

    def __lowercase   (   self    :  Any		)   ->  Optional[int]:
      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 __lowercase   (   self    :  Any		)   ->  str:
      pass

    @unittest.skip(reason="""Cvt does not use inputs_embeds"""		)
    def __lowercase   (   self    :  str		)   ->  int:
      pass

    @unittest.skip(reason="""Cvt does not support input and output embeddings"""		)
    def __lowercase   (   self    :  Optional[Any]		)   ->  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 __lowercase   (   self    :  List[str]		)   ->  Optional[Any]:
      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 __lowercase   (   self    :  Optional[int]		)   ->  Dict:
      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 __lowercase   (   self    :  Union[str, Any]		)   ->  Tuple:
      lowerCAmelCase_			:      List[Any]			=   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 __lowercase   (   self    :  Optional[Any]		)   ->  Union[str, Any]:
      lowerCAmelCase_			:      str			=   self.model_tester.prepare_config_and_inputs_for_common()

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

        lowerCAmelCase_			:      Any			=   ['''pixel_values''']
        self.assertListEqual(arg_names[:1]    ,						_A		)

    def __lowercase   (   self    :  Any		)   ->  Optional[Any]:


      def check_hidden_states_output(lowerCamelCase    :  Dict    ,						lowerCamelCase    :  Optional[Any]    ,						lowerCamelCase    :  Dict		):
        lowerCAmelCase_			:      str			=   model_class(_A		)

        lowerCAmelCase_			:      List[str]			=   model(**self._prepare_for_class(_A    ,						_A		)		)
        lowerCAmelCase_			:      Tuple			=   outputs.hidden_states

        lowerCAmelCase_			:      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,
            ]    ,						)

      lowerCAmelCase_			:      int			=   self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        lowerCAmelCase_			:      Tuple			=   True
        check_hidden_states_output(_A    ,						_A    ,						_A		)

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        lowerCAmelCase_			:      List[str]			=   True

        check_hidden_states_output(_A    ,						_A    ,						_A		)

    def __lowercase   (   self    :  Tuple		)   ->  Dict:
      lowerCAmelCase_			:      Union[str, Any]			=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*_A		)

    def __lowercase   (   self    :  str		)   ->  List[str]:
      lowerCAmelCase_			:      Optional[Any]			=   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*_A		)





    @slow
    def __lowercase   (   self    :  Optional[Any]		)   ->  Union[str, Any]:


      for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        lowerCAmelCase_			:      Optional[int]			=   TFCvtModel.from_pretrained(_A		)
        self.assertIsNotNone(_A		)







def 				UpperCamelCase_			(    ):
  '''simple docstring'''





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






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







    """simple docstring"""



    @cached_property
    def __lowercase   (   self    :  Union[str, Any]		)   ->  str:
      return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]		)





    @slow
    def __lowercase   (   self    :  Any		)   ->  int:
      lowerCAmelCase_			:      Dict			=   TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]		)

      lowerCAmelCase_			:      Union[str, Any]			=   self.default_image_processor
      lowerCAmelCase_			:      Optional[Any]			=   prepare_img()
      lowerCAmelCase_			:      Tuple			=   image_processor(images=_A    ,						return_tensors="""tf"""		)

      # forward pass
      lowerCAmelCase_			:      Optional[Any]			=   model(**_A		)

      # verify the logits
      lowerCAmelCase_			:      Union[str, Any]			=   tf.TensorShape((1, 10_00)		)
      self.assertEqual(outputs.logits.shape    ,						_A		)

      lowerCAmelCase_			:      Union[str, Any]			=   tf.constant([0.9_285, 0.9_015, -0.3_150]		)
      self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()    ,						_A    ,						atol=1E-4		)		)







'''simple docstring'''
import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


UpperCamelCase__    					=      logging.get_logger(__name__)





def 					a__       (					lowerCAmelCase__						)    ->	Optional[Any]:
    UpperCAmelCase__			:		str													=				R'''\w+[.]\d+'''
    UpperCAmelCase__			:		List[Any]													=				re.findall(lowerCAmelCase__  ,					lowerCAmelCase__						)
    for pat in pats:
        UpperCAmelCase__			:		Union[str, Any]													=				key.replace(lowerCAmelCase__  ,					'''_'''.join(pat.split('''.'''						)						)						)
    return key





def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)    ->	Union[str, Any]:
    UpperCAmelCase__			:		Optional[Any]													=				pt_tuple_key[:-1] + ('''scale''',)
    if (
        any('''norm''' in str_ for str_ in pt_tuple_key						)
        and (pt_tuple_key[-1] == "bias")
        and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
        and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
    ):
        UpperCAmelCase__			:		Union[str, Any]													=				pt_tuple_key[:-1] + ('''scale''',)
        return renamed_pt_tuple_key, pt_tensor
    elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
        UpperCAmelCase__			:		Optional[int]													=				pt_tuple_key[:-1] + ('''scale''',)
        return renamed_pt_tuple_key, pt_tensor

    # embedding
    if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
        UpperCAmelCase__			:		str													=				pt_tuple_key[:-1] + ('''embedding''',)
        return renamed_pt_tuple_key, pt_tensor

    # conv layer
    UpperCAmelCase__			:		Optional[Any]													=				pt_tuple_key[:-1] + ('''kernel''',)
    if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
        UpperCAmelCase__			:		List[Any]													=				pt_tensor.transpose(2  ,					3  ,					1  ,					0						)
        return renamed_pt_tuple_key, pt_tensor

    # linear layer
    UpperCAmelCase__			:		int													=				pt_tuple_key[:-1] + ('''kernel''',)
    if pt_tuple_key[-1] == "weight":
        UpperCAmelCase__			:		Optional[Any]													=				pt_tensor.T
        return renamed_pt_tuple_key, pt_tensor

    # old PyTorch layer norm weight
    UpperCAmelCase__			:		Union[str, Any]													=				pt_tuple_key[:-1] + ('''weight''',)
    if pt_tuple_key[-1] == "gamma":
        return renamed_pt_tuple_key, pt_tensor

    # old PyTorch layer norm bias
    UpperCAmelCase__			:		Optional[Any]													=				pt_tuple_key[:-1] + ('''bias''',)
    if pt_tuple_key[-1] == "beta":
        return renamed_pt_tuple_key, pt_tensor

    return pt_tuple_key, pt_tensor







def 					a__       (					lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__=42						)    ->	Tuple:
    # Step 1: Convert pytorch tensor to numpy
    UpperCAmelCase__			:		int													=				{k: v.numpy() for k, v in pt_state_dict.items()}

    # Step 2: Since the model is stateless, get random Flax params
    UpperCAmelCase__			:		Tuple													=				flax_model.init_weights(PRNGKey(lowerCAmelCase__						)						)

    UpperCAmelCase__			:		Optional[Any]													=				flatten_dict(lowerCAmelCase__						)
    UpperCAmelCase__			:		List[str]													=				{}

    # Need to change some parameters name to match Flax names
    for pt_key, pt_tensor in pt_state_dict.items():
        UpperCAmelCase__			:		Optional[int]													=				rename_key(lowerCAmelCase__						)
        UpperCAmelCase__			:		str													=				tuple(renamed_pt_key.split('''.'''						)						)

        # Correctly rename weight parameters
        UpperCAmelCase__						,						UpperCAmelCase__			:		Optional[int]													=				rename_key_and_reshape_tensor(lowerCAmelCase__  ,					lowerCAmelCase__  ,					lowerCAmelCase__						)

        if flax_key in random_flax_state_dict:
            if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                raise ValueError(
                    F"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
                    F"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""						)

        # also add unexpected weight so that warning is thrown
        UpperCAmelCase__			:		List[str]													=				jnp.asarray(lowerCAmelCase__						)

    return unflatten_dict(lowerCAmelCase__						)
