




"""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     __magic_name__							(       UpperCAmelCase__				):






				'''simple docstring'''





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.config_class(**self.inputs_dict )
							self.parent.assertTrue(hasattr(_a		, """embed_dim""" ) )
							self.parent.assertTrue(hasattr(_a		, """num_heads""" ) )






class     __magic_name__							:






				'''simple docstring'''





				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-1_2		, _a=True		, _a=True		, _a=2		, ):





							"""simple docstring"""
							lowerCamelCase			      =					parent
							lowerCamelCase			      =					batch_size
							lowerCamelCase			      =					image_size
							lowerCamelCase			      =					patch_sizes
							lowerCamelCase			      =					patch_stride
							lowerCamelCase			      =					patch_padding
							lowerCamelCase			      =					is_training
							lowerCamelCase			      =					use_labels
							lowerCamelCase			      =					num_labels
							lowerCamelCase			      =					num_channels
							lowerCamelCase			      =					embed_dim
							lowerCamelCase			      =					num_heads
							lowerCamelCase			      =					stride_kv
							lowerCamelCase			      =					depth
							lowerCamelCase			      =					cls_token
							lowerCamelCase			      =					attention_drop_rate
							lowerCamelCase			      =					initializer_range
							lowerCamelCase			      =					layer_norm_eps





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

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

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





				def    _lowerCAmelCase				(       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    _lowerCAmelCase				(       self		, _a		, _a		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					TFCvtModel(config=_a )
							lowerCamelCase			      =					model(_a		, training=_a )
							lowerCamelCase			      =					(self.image_size, self.image_size)
							lowerCamelCase      ,   lowerCamelCase			      =					image_size[0], image_size[1]
							for i in range(len(self.depth ) ):
										lowerCamelCase			      =					floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1 )
										lowerCamelCase			      =					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    _lowerCAmelCase				(       self		, _a		, _a		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					self.num_labels
							lowerCamelCase			      =					TFCvtForImageClassification(_a )
							lowerCamelCase			      =					model(_a		, labels=_a		, training=_a )
							self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.num_labels) )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.prepare_config_and_inputs()
							lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					config_and_inputs
							lowerCamelCase			      =					{"""pixel_values""": pixel_values}
							return config, inputs_dict






@require_tf
class     __magic_name__							(       UpperCAmelCase__						,					UpperCAmelCase__						,					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    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					TFCvtModelTester(self )
							lowerCamelCase			      =					TFCvtConfigTester(self		, config_class=_a		, has_text_modality=_a		, hidden_size=37 )





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





							"""simple docstring"""
							pass





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





							"""simple docstring"""
							pass





				@unittest.skip(reason="""Cvt does not support input and output embeddings""" )
				def    _lowerCAmelCase				(       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    _lowerCAmelCase				(       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    _lowerCAmelCase				(       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    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					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    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase      ,   lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs_for_common()

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

										lowerCamelCase			      =					["""pixel_values"""]
										self.assertListEqual(arg_names[:1]		, _a )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							def check_hidden_states_output(_a		, _a		, _a ):
										lowerCamelCase			      =					model_class(_a )

										lowerCamelCase			      =					model(**self._prepare_for_class(_a		, _a ) )
										lowerCamelCase			      =					outputs.hidden_states

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

							for model_class in self.all_model_classes:
										lowerCamelCase			      =					True
										check_hidden_states_output(_a		, _a		, _a )

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

										check_hidden_states_output(_a		, _a		, _a )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_a )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*_a )





				@slow
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCamelCase			      =					TFCvtModel.from_pretrained(_a )
										self.assertIsNotNone(_a )




def     a__ (				)						->					Optional[int]:
			lowerCamelCase			      =					Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""				)
			return image






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






				'''simple docstring'''





				@cached_property
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0] )





				@slow
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0] )

							lowerCamelCase			      =					self.default_image_processor
							lowerCamelCase			      =					prepare_img()
							lowerCamelCase			      =					image_processor(images=_a		, return_tensors="""tf""" )

							# forward pass
							lowerCamelCase			      =					model(**_a )

							# verify the logits
							lowerCamelCase			      =					tf.TensorShape((1, 1_000) )
							self.assertEqual(outputs.logits.shape		, _a )

							lowerCamelCase			      =					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"""




from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


lowerCAmelCase							:			Tuple     							=							logging.get_logger(__name__)



def     a__ (				snake_case__       , snake_case__				)						->					Tuple:
			try:
						with open(snake_case__       , """rb"""				) as flax_state_f:
									lowerCamelCase			      =					from_bytes(snake_case__       , flax_state_f.read()				)
			except UnpicklingError as e:
						try:
									with open(snake_case__				) as f:
												if f.read().startswith("""version"""				):
															raise OSError(
															    """You seem to have cloned a repository without having git-lfs installed. Please"""
															    """ install git-lfs and run `git lfs install` followed by `git lfs pull` in the"""
															    """ folder you cloned."""				)
												else:
															raise ValueError from e
						except (UnicodeDecodeError, ValueError):
									raise EnvironmentError(F'Unable to convert {model_file} to Flax deserializable object. '				)

			return load_flax_weights_in_pytorch_model(snake_case__       , snake_case__				)



def     a__ (				snake_case__       , snake_case__				)						->					Tuple:

			try:
						import torch  # noqa: F401
			except ImportError:
						logger.error(
						    """Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see"""
						    """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
						    """ instructions."""				)
						raise

			# check if we have bf16 weights
			lowerCamelCase			      =					flatten_dict(jax.tree_util.tree_map(lambda snake_case__				: x.dtype == jnp.bfloataa       , snake_case__				)				).values()
			if any(snake_case__				):
						# convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

						# and bf16 is not fully supported in PT yet.
						logger.warning(
						    """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
						    """before loading those in PyTorch model."""				)
						lowerCamelCase			      =					jax.tree_util.tree_map(
						    lambda snake_case__				: params.astype(np.floataa				) if params.dtype == jnp.bfloataa else params       , snake_case__				)

			lowerCamelCase			      =					""""""

			lowerCamelCase			      =					flatten_dict(snake_case__       , sep="""."""				)
			lowerCamelCase			      =					pt_model.state_dict()

			# keep track of unexpected & missing keys
			lowerCamelCase			      =					[]
			lowerCamelCase			      =					set(pt_model_dict.keys()				)

			for flax_key_tuple, flax_tensor in flax_state_dict.items():
						lowerCamelCase			      =					flax_key_tuple.split("""."""				)

						if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
									lowerCamelCase			      =					flax_key_tuple_array[:-1] + ["""weight"""]
									lowerCamelCase			      =					jnp.transpose(snake_case__       , (3, 2, 0, 1)				)
						elif flax_key_tuple_array[-1] == "kernel":
									lowerCamelCase			      =					flax_key_tuple_array[:-1] + ["""weight"""]
									lowerCamelCase			      =					flax_tensor.T
						elif flax_key_tuple_array[-1] == "scale":
									lowerCamelCase			      =					flax_key_tuple_array[:-1] + ["""weight"""]

						if "time_embedding" not in flax_key_tuple_array:
									for i, flax_key_tuple_string in enumerate(snake_case__				):
												lowerCamelCase			      =					(
												    flax_key_tuple_string.replace("""_0"""       , """.0"""				)
												    .replace("""_1"""       , """.1"""				)
												    .replace("""_2"""       , """.2"""				)
												    .replace("""_3"""       , """.3"""				)
												    .replace("""_4"""       , """.4"""				)
												    .replace("""_5"""       , """.5"""				)
												    .replace("""_6"""       , """.6"""				)
												    .replace("""_7"""       , """.7"""				)
												    .replace("""_8"""       , """.8"""				)
												    .replace("""_9"""       , """.9"""				)
												)

						lowerCamelCase			      =					""".""".join(snake_case__				)

						if flax_key in pt_model_dict:
									if flax_tensor.shape != pt_model_dict[flax_key].shape:
												raise ValueError(
												    F'Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '
												    F'to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'				)
									else:
												# add weight to pytorch dict
												lowerCamelCase			      =					np.asarray(snake_case__				) if not isinstance(snake_case__       , np.ndarray				) else flax_tensor
												lowerCamelCase			      =					torch.from_numpy(snake_case__				)
												# remove from missing keys
												missing_keys.remove(snake_case__				)
						else:
									# weight is not expected by PyTorch model
									unexpected_keys.append(snake_case__				)

			pt_model.load_state_dict(snake_case__				)

			# re-transform missing_keys to list
			lowerCamelCase			      =					list(snake_case__				)

			if len(snake_case__				) > 0:
						logger.warning(
						    """Some weights of the Flax model were not used when initializing the PyTorch model"""
						    F' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'
						    F' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'
						    """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
						    F' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'
						    """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
						    """ FlaxBertForSequenceClassification model)."""				)
			if len(snake_case__				) > 0:
						logger.warning(
						    F'Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'
						    F' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'
						    """ use it for predictions and inference."""				)

			return pt_model



