



"""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 				A_      (					SCREAMING_SNAKE_CASE_    ):



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:int )    ->      List[str]:
												UpperCAmelCase					=			self.config_class(**self.inputs_dict )
												self.parent.assertTrue(hasattr(lowercase_       ,  'embed_dim' ) )
												self.parent.assertTrue(hasattr(lowercase_       ,  'num_heads' ) )







class 				A_      :



						"""simple docstring"""






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






						def      UpperCAmelCase__					( self			:Union[str, Any] )    ->      List[Any]:
												UpperCAmelCase					=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

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

												UpperCAmelCase					=			self.get_config()
												return config, pixel_values, labels






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






						def      UpperCAmelCase__					( self			:List[str] )    ->      str:
												UpperCAmelCase					=			self.prepare_config_and_inputs()
												UpperCAmelCase				,   UpperCAmelCase				,   UpperCAmelCase					=			config_and_inputs
												UpperCAmelCase					=			{'pixel_values': pixel_values}
												return config, inputs_dict







@require_tf
class 				A_      (					SCREAMING_SNAKE_CASE_     ,       SCREAMING_SNAKE_CASE_     ,       unittest.TestCase    ):



						"""simple docstring"""
						__UpperCamelCase		    =							(TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
						__UpperCamelCase		    =							(
						    {"""feature-extraction""": TFCvtModel, """image-classification""": TFCvtForImageClassification}
						    if is_tf_available()
						    else {}
						)
						__UpperCamelCase		    =							False
						__UpperCamelCase		    =							False
						__UpperCamelCase		    =							False
						__UpperCamelCase		    =							False
						__UpperCamelCase		    =							False






						def      UpperCAmelCase__					( self			:Any )    ->      Optional[Any]:
												UpperCAmelCase					=			TFCvtModelTester(self )
												UpperCAmelCase					=			TFCvtConfigTester(self       ,  config_class=lowercase_       ,  has_text_modality=lowercase_       ,  hidden_size=37 )






						def      UpperCAmelCase__					( self			:List[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      UpperCAmelCase__					( self			:Dict )    ->      Any:
												pass






						@unittest.skip(reason='Cvt does not use inputs_embeds' )
						def      UpperCAmelCase__					( self			:str )    ->      Dict:
												pass






						@unittest.skip(reason='Cvt does not support input and output embeddings' )
						def      UpperCAmelCase__					( self			:List[Any] )    ->      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      UpperCAmelCase__					( self			:Optional[int] )    ->      Tuple:
												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			: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      UpperCAmelCase__					( self			:Dict )    ->      Optional[int]:
												UpperCAmelCase					=			tf.keras.mixed_precision.Policy('mixed_float16' )
												tf.keras.mixed_precision.set_global_policy(lowercase_ )
												super().test_keras_fit()
												tf.keras.mixed_precision.set_global_policy('float32' )






						def      UpperCAmelCase__					( self			:Dict )    ->      Union[str, Any]:
												UpperCAmelCase				,   UpperCAmelCase					=			self.model_tester.prepare_config_and_inputs_for_common()

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

																		UpperCAmelCase					=			['pixel_values']
																		self.assertListEqual(arg_names[:1]       ,  lowercase_ )






						def      UpperCAmelCase__					( self			:Tuple )    ->      Optional[int]:
												def check_hidden_states_output(lowercase_			:Union[str, Any]       ,  lowercase_			:str       ,  lowercase_			:List[str] ):
																		UpperCAmelCase					=			model_class(lowercase_ )

																		UpperCAmelCase					=			model(**self._prepare_for_class(lowercase_       ,  lowercase_ ) )
																		UpperCAmelCase					=			outputs.hidden_states

																		UpperCAmelCase					=			len(self.model_tester.depth )
																		self.assertEqual(len(lowercase_ )       ,  lowercase_ )

																		# 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					=			self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																		UpperCAmelCase					=			True
																		check_hidden_states_output(lowercase_       ,  lowercase_       ,  lowercase_ )

																		# check that output_hidden_states also work using config
																		del inputs_dict["output_hidden_states"]
																		UpperCAmelCase					=			True

																		check_hidden_states_output(lowercase_       ,  lowercase_       ,  lowercase_ )






						def      UpperCAmelCase__					( self			:int )    ->      str:
												UpperCAmelCase					=			self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*lowercase_ )






						def      UpperCAmelCase__					( self			:List[Any] )    ->      List[Any]:
												UpperCAmelCase					=			self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_image_classification(*lowercase_ )






						@slow
						def      UpperCAmelCase__					( self			:str )    ->      Optional[Any]:
												for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		UpperCAmelCase					=			TFCvtModel.from_pretrained(lowercase_ )
																		self.assertIsNotNone(lowercase_ )


def 			_lowerCAmelCase				(       ):
						UpperCAmelCase					=			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'						)
						return image







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



						"""simple docstring"""






						@cached_property
						def      UpperCAmelCase__					( self			:Dict )    ->      str:
												return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0] )






						@slow
						def      UpperCAmelCase__					( self			:List[str] )    ->      Tuple:
												UpperCAmelCase					=			TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0] )

												UpperCAmelCase					=			self.default_image_processor
												UpperCAmelCase					=			prepare_img()
												UpperCAmelCase					=			image_processor(images=lowercase_       ,  return_tensors='tf' )

												# forward pass
												UpperCAmelCase					=			model(**lowercase_ )

												# verify the logits
												UpperCAmelCase					=			tf.TensorShape((1, 10_00) )
												self.assertEqual(outputs.logits.shape       ,  lowercase_ )

												UpperCAmelCase					=			tf.constant([0.9285, 0.9015, -0.3150] )
												self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()       ,  lowercase_       ,  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


snake_case_    					=	logging.get_logger(__name__)

def 			_lowerCAmelCase				(       lowercase_						):
						UpperCAmelCase					=			R'\w+[.]\d+'
						UpperCAmelCase					=			re.findall(lowercase_		,     lowercase_						)
						for pat in pats:
												UpperCAmelCase					=			key.replace(lowercase_		,     '_'.join(pat.split('.'						)						)						)
						return key

def 			_lowerCAmelCase				(       lowercase_		,     lowercase_		,     lowercase_						):
						UpperCAmelCase					=			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					=			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					=			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					=			pt_tuple_key[:-1] + ('embedding',)
												return renamed_pt_tuple_key, pt_tensor

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

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

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

						# old PyTorch layer norm bias
						UpperCAmelCase					=			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 			_lowerCAmelCase				(       lowercase_		,     lowercase_		,     lowercase_=42						):
						# Step 1: Convert pytorch tensor to numpy
						UpperCAmelCase					=			{k: v.numpy() for k, v in pt_state_dict.items()}

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

						UpperCAmelCase					=			flatten_dict(lowercase_						)
						UpperCAmelCase					=			{}

						# Need to change some parameters name to match Flax names
						for pt_key, pt_tensor in pt_state_dict.items():
												UpperCAmelCase					=			rename_key(lowercase_						)
												UpperCAmelCase					=			tuple(renamed_pt_key.split('.'						)						)

												# Correctly rename weight parameters
												UpperCAmelCase				,   UpperCAmelCase					=			rename_key_and_reshape_tensor(lowercase_		,     lowercase_		,     lowercase_						)

												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					=			jnp.asarray(lowercase_						)

						return unflatten_dict(lowercase_						)





