


"""simple docstring"""




from typing import Optional, Union

import torch
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithPoolingAndNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_mobilenet_va import MobileNetVaConfig


lowercase_			     =  logging.get_logger(__name__)


# General docstring
lowercase_			     =  "MobileNetV1Config"

# Base docstring
lowercase_			     =  "google/mobilenet_v1_1.0_224"
lowercase_			     =  [1, 1_0_2_4, 7, 7]

# Image classification docstring
lowercase_			     =  "google/mobilenet_v1_1.0_224"
lowercase_			     =  "tabby, tabby cat"


lowercase_			     =  [
    "google/mobilenet_v1_1.0_224",
    "google/mobilenet_v1_0.75_192",
    # See all MobileNetV1 models at https://huggingface.co/models?filter=mobilenet_v1
]

def        lowercase							(		lowerCAmelCase__ :     List[str]    ,					lowerCAmelCase__ :     int    ,					lowerCAmelCase__ :     Optional[int]=None			) ->      int:
		__a           =			{}

		if isinstance(lowerCAmelCase__    ,					lowerCAmelCase__			):
				__a           =			model.mobilenet_va
		else:
				__a           =			model

		__a           =			'''MobilenetV1/Conv2d_0/'''
		__a           =			backbone.conv_stem.convolution.weight
		__a           =			backbone.conv_stem.normalization.bias
		__a           =			backbone.conv_stem.normalization.weight
		__a           =			backbone.conv_stem.normalization.running_mean
		__a           =			backbone.conv_stem.normalization.running_var

		for i in range(13			):
				__a           =			i + 1
				__a           =			i * 2

				__a           =			backbone.layer[pt_index]
				__a           =			f'''MobilenetV1/Conv2d_{tf_index}_depthwise/'''
				__a           =			pointer.convolution.weight
				__a           =			pointer.normalization.bias
				__a           =			pointer.normalization.weight
				__a           =			pointer.normalization.running_mean
				__a           =			pointer.normalization.running_var

				__a           =			backbone.layer[pt_index + 1]
				__a           =			f'''MobilenetV1/Conv2d_{tf_index}_pointwise/'''
				__a           =			pointer.convolution.weight
				__a           =			pointer.normalization.bias
				__a           =			pointer.normalization.weight
				__a           =			pointer.normalization.running_mean
				__a           =			pointer.normalization.running_var

		if isinstance(lowerCAmelCase__    ,					lowerCAmelCase__			):
				__a           =			'''MobilenetV1/Logits/Conv2d_1c_1x1/'''
				__a           =			model.classifier.weight
				__a           =			model.classifier.bias

		return tf_to_pt_map

def        lowercase							(		lowerCAmelCase__ :     Union[str, Any]    ,					lowerCAmelCase__ :     List[str]    ,					lowerCAmelCase__ :     int			) ->      Tuple:
		try:
				import numpy as np
				import tensorflow as tf
		except ImportError:
				logger.error(
				    '''Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see '''
				    '''https://www.tensorflow.org/install/ for installation instructions.'''			)
				raise

		# Load weights from TF model
		__a           =			tf.train.list_variables(lowerCAmelCase__			)
		__a           =			{}
		for name, shape in init_vars:
				logger.info(f'''Loading TF weight {name} with shape {shape}'''			)
				__a           =			tf.train.load_variable(lowerCAmelCase__    ,					lowerCAmelCase__			)
				__a           =			array

		# Build TF to PyTorch weights loading map
		__a           =			_build_tf_to_pytorch_map(lowerCAmelCase__    ,					lowerCAmelCase__    ,					lowerCAmelCase__			)

		for name, pointer in tf_to_pt_map.items():
				logger.info(f'''Importing {name}'''			)
				if name not in tf_weights:
						logger.info(f'''{name} not in tf pre-trained weights, skipping'''			)
						continue

				__a           =			tf_weights[name]

				if "depthwise_weights" in name:
						logger.info('''Transposing depthwise'''			)
						__a           =			np.transpose(lowerCAmelCase__    ,					(2, 3, 0, 1)			)
				elif "weights" in name:
						logger.info('''Transposing'''			)
						if len(pointer.shape			) == 2:  # copying into linear layer
								__a           =			array.squeeze().transpose()
						else:
								__a           =			np.transpose(lowerCAmelCase__    ,					(3, 2, 0, 1)			)

				if pointer.shape != array.shape:
						raise ValueError(f'''Pointer shape {pointer.shape} and array shape {array.shape} mismatched'''			)

				logger.info(f'''Initialize PyTorch weight {name} {array.shape}'''			)
				__a           =			torch.from_numpy(lowerCAmelCase__			)

				tf_weights.pop(lowerCAmelCase__    ,					lowerCAmelCase__			)
				tf_weights.pop(name + '''/RMSProp'''    ,					lowerCAmelCase__			)
				tf_weights.pop(name + '''/RMSProp_1'''    ,					lowerCAmelCase__			)
				tf_weights.pop(name + '''/ExponentialMovingAverage'''    ,					lowerCAmelCase__			)

		logger.info(f'''Weights not copied to PyTorch model: {', '.join(tf_weights.keys()			)}'''			)
		return model

def        lowercase							(		lowerCAmelCase__ :     torch.Tensor    ,					lowerCAmelCase__ :     nn.Convad			) ->      torch.Tensor:
		__a      ,     __a           =			features.shape[-2:]
		__a      ,     __a           =			conv_layer.stride
		__a      ,     __a           =			conv_layer.kernel_size

		if in_height % stride_height == 0:
				__a           =			max(kernel_height - stride_height    ,					0			)
		else:
				__a           =			max(kernel_height - (in_height % stride_height)    ,					0			)

		if in_width % stride_width == 0:
				__a           =			max(kernel_width - stride_width    ,					0			)
		else:
				__a           =			max(kernel_width - (in_width % stride_width)    ,					0			)

		__a           =			pad_along_width // 2
		__a           =			pad_along_width - pad_left
		__a           =			pad_along_height // 2
		__a           =			pad_along_height - pad_top

		__a           =			(pad_left, pad_right, pad_top, pad_bottom)
		return nn.functional.pad(lowerCAmelCase__    ,					lowerCAmelCase__    ,					'''constant'''    ,					0.0			)
class 					__lowerCAmelCase					(   nn.Module       ):


		'''simple docstring'''



		def __init__(  self							, _a							, _a							, _a							, _a							, _a = 1							, _a = 1							, _a = False							, _a = True							, _a = True							, ):
				super().__init__()
				__a           =			config

				if in_channels % groups != 0:
						raise ValueError(f'''Input channels ({in_channels}) are not divisible by {groups} groups.'''      )
				if out_channels % groups != 0:
						raise ValueError(f'''Output channels ({out_channels}) are not divisible by {groups} groups.'''      )

				__a           =			0 if config.tf_padding else int((kernel_size - 1) / 2      )

				__a           =			nn.Convad(
				    in_channels=_a							, out_channels=_a							, kernel_size=_a							, stride=_a							, padding=_a							, groups=_a							, bias=_a							, padding_mode='''zeros'''							, )

				if use_normalization:
						__a           =			nn.BatchNormad(
						    num_features=_a							, eps=config.layer_norm_eps							, momentum=0.9997							, affine=_a							, track_running_stats=_a							, )
				else:
						__a           =			None

				if use_activation:
						if isinstance(_a							, _a      ):
								__a           =			ACTaFN[use_activation]
						elif isinstance(config.hidden_act							, _a      ):
								__a           =			ACTaFN[config.hidden_act]
						else:
								__a           =			config.hidden_act
				else:
						__a           =			None





		def 			__UpperCAmelCase       (  self							, _a      ):
				if self.config.tf_padding:
						__a           =			apply_tf_padding(_a							, self.convolution      )
				__a           =			self.convolution(_a      )
				if self.normalization is not None:
						__a           =			self.normalization(_a      )
				if self.activation is not None:
						__a           =			self.activation(_a      )
				return features







class 					__lowerCAmelCase					(   __SCREAMING_SNAKE_CASE       ):


		'''simple docstring'''



		__UpperCAmelCase		:			str										=				MobileNetVaConfig
		__UpperCAmelCase		:			Optional[int]										=				load_tf_weights_in_mobilenet_va
		__UpperCAmelCase		:			Optional[Any]										=				'mobilenet_v1'
		__UpperCAmelCase		:			Tuple										=				'pixel_values'
		__UpperCAmelCase		:			int										=				False

		def 			__UpperCAmelCase       (  self							, _a      ):
				if isinstance(_a							, (nn.Linear, nn.Convad)      ):
						module.weight.data.normal_(mean=0.0							, std=self.config.initializer_range      )
						if module.bias is not None:
								module.bias.data.zero_()
				elif isinstance(_a							, nn.BatchNormad      ):
						module.bias.data.zero_()
						module.weight.data.fill_(1.0      )







lowercase_			     =  r"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`MobileNetV1Config`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

lowercase_			     =  r"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`MobileNetV1ImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"
@add_start_docstrings(
    'The bare MobileNetV1 model outputting raw hidden-states without any specific head on top.'     ,			__SCREAMING_SNAKE_CASE     ,			)
class 					__lowerCAmelCase					(   __SCREAMING_SNAKE_CASE       ):


		'''simple docstring'''



		def __init__(  self							, _a							, _a = True      ):
				super().__init__(_a      )
				__a           =			config

				__a           =			32
				__a           =			max(int(depth * config.depth_multiplier      )							, config.min_depth      )

				__a           =			MobileNetVaConvLayer(
				    _a							, in_channels=config.num_channels							, out_channels=_a							, kernel_size=3							, stride=2							, )

				__a           =			[1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

				__a           =			nn.ModuleList()
				for i in range(13      ):
						__a           =			out_channels

						if strides[i] == 2 or i == 0:
								depth *= 2
								__a           =			max(int(depth * config.depth_multiplier      )							, config.min_depth      )

						self.layer.append(
						    MobileNetVaConvLayer(
						        _a							, in_channels=_a							, out_channels=_a							, kernel_size=3							, stride=strides[i]							, groups=_a							, )      )

						self.layer.append(
						    MobileNetVaConvLayer(
						        _a							, in_channels=_a							, out_channels=_a							, kernel_size=1							, )      )

				__a           =			nn.AdaptiveAvgPoolad((1, 1)      ) if add_pooling_layer else None

				# Initialize weights and apply final processing
				self.post_init()

		def 			__UpperCAmelCase       (  self							, _a      ):
				raise NotImplementedError





		@add_start_docstrings_to_model_forward(_a      )
		@add_code_sample_docstrings(
		    checkpoint=_CHECKPOINT_FOR_DOC							, output_type=_a							, config_class=_CONFIG_FOR_DOC							, modality='''vision'''							, expected_output=_EXPECTED_OUTPUT_SHAPE							, )
		def 			__UpperCAmelCase       (  self							, _a = None							, _a = None							, _a = None							, ):
				__a           =			(
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				__a           =			return_dict if return_dict is not None else self.config.use_return_dict

				if pixel_values is None:
						raise ValueError('''You have to specify pixel_values'''      )

				__a           =			self.conv_stem(_a      )

				__a           =			() if output_hidden_states else None

				for i, layer_module in enumerate(self.layer      ):
						__a           =			layer_module(_a      )

						if output_hidden_states:
								__a           =			all_hidden_states + (hidden_states,)

				__a           =			hidden_states

				if self.pooler is not None:
						__a           =			torch.flatten(self.pooler(_a      )							, start_dim=1      )
				else:
						__a           =			None

				if not return_dict:
						return tuple(v for v in [last_hidden_state, pooled_output, all_hidden_states] if v is not None      )

				return BaseModelOutputWithPoolingAndNoAttention(
				    last_hidden_state=_a							, pooler_output=_a							, hidden_states=_a							, )







@add_start_docstrings(
    '\n    MobileNetV1 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    '     ,			__SCREAMING_SNAKE_CASE     ,			)
class 					__lowerCAmelCase					(   __SCREAMING_SNAKE_CASE       ):


		'''simple docstring'''



		def __init__(  self							, _a      ):
				super().__init__(_a      )

				__a           =			config.num_labels
				__a           =			MobileNetVaModel(_a      )

				__a           =			self.mobilenet_va.layer[-1].convolution.out_channels

				# Classifier head
				__a           =			nn.Dropout(config.classifier_dropout_prob							, inplace=_a      )
				__a           =			nn.Linear(_a							, config.num_labels      ) if config.num_labels > 0 else nn.Identity()

				# Initialize weights and apply final processing
				self.post_init()





		@add_start_docstrings_to_model_forward(_a      )
		@add_code_sample_docstrings(
		    checkpoint=_IMAGE_CLASS_CHECKPOINT							, output_type=_a							, config_class=_CONFIG_FOR_DOC							, expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							, )
		def 			__UpperCAmelCase       (  self							, _a = None							, _a = None							, _a = None							, _a = None							, ):
				__a           =			return_dict if return_dict is not None else self.config.use_return_dict

				__a           =			self.mobilenet_va(_a							, output_hidden_states=_a							, return_dict=_a      )

				__a           =			outputs.pooler_output if return_dict else outputs[1]

				__a           =			self.classifier(self.dropout(_a      )      )

				__a           =			None
				if labels is not None:
						if self.config.problem_type is None:
								if self.num_labels == 1:
										__a           =			'''regression'''
								elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
										__a           =			'''single_label_classification'''
								else:
										__a           =			'''multi_label_classification'''

						if self.config.problem_type == "regression":
								__a           =			MSELoss()
								if self.num_labels == 1:
										__a           =			loss_fct(logits.squeeze()							, labels.squeeze()      )
								else:
										__a           =			loss_fct(_a							, _a      )
						elif self.config.problem_type == "single_label_classification":
								__a           =			CrossEntropyLoss()
								__a           =			loss_fct(logits.view(-1							, self.num_labels      )							, labels.view(-1      )      )
						elif self.config.problem_type == "multi_label_classification":
								__a           =			BCEWithLogitsLoss()
								__a           =			loss_fct(_a							, _a      )

				if not return_dict:
						__a           =			(logits,) + outputs[2:]
						return ((loss,) + output) if loss is not None else output

				return ImageClassifierOutputWithNoAttention(
				    loss=_a							, logits=_a							, hidden_states=outputs.hidden_states							, )











"""simple docstring"""






import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_flax_cross_test,
    require_flax,
    require_torch,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import is_flax_available, is_torch_available, is_vision_available

from ...test_modeling_flax_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_flax_bert import FlaxBertModelTester
from ..clip.test_modeling_flax_clip import FlaxCLIPVisionModelTester
from ..vit.test_modeling_flax_vit import FlaxViTModelTester


if is_flax_available():
	from transformers import (
	    FlaxBertModel,
	    FlaxCLIPVisionModel,
	    FlaxVisionTextDualEncoderModel,
	    FlaxViTModel,
	    VisionTextDualEncoderConfig,
	    VisionTextDualEncoderProcessor,
	)
	from transformers.modeling_flax_pytorch_utils import (
	    convert_pytorch_state_dict_to_flax,
	    load_flax_weights_in_pytorch_model,
	)


if is_torch_available():
	import torch

	from transformers import VisionTextDualEncoderModel

if is_vision_available():
	from PIL import Image


def 			_lowercase				(  __snake_case	)						->   List[str]:
			if isinstance(__snake_case					,collections.abc.Iterable	):
						return x
			return (x, x)



@require_flax
class 			A__   :




	'''simple docstring'''





	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		Dict       ,			_SCREAMING_SNAKE_CASE:		Optional[Any])				->       str:

				"""simple docstring"""

				pass

	def 							_SCREAMING_SNAKE_CASE   (     self:		str)				->       int:

				"""simple docstring"""

				pass

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str])				->       Tuple:

				"""simple docstring"""

				pass

	def 							_SCREAMING_SNAKE_CASE   (     self:		Any       ,			_SCREAMING_SNAKE_CASE:		np.ndarray       ,			_SCREAMING_SNAKE_CASE:		np.ndarray       ,			_SCREAMING_SNAKE_CASE:		float)				->       List[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  str												=       np.abs((a - b)).max()
				self.assertLessEqual(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			F"""Difference between torch and flax is {diff} (>= {tol}).""")

	def 							_SCREAMING_SNAKE_CASE   (     self:		str       ,			_SCREAMING_SNAKE_CASE:		List[str]       ,			_SCREAMING_SNAKE_CASE:		Dict       ,			_SCREAMING_SNAKE_CASE:		Any       ,			_SCREAMING_SNAKE_CASE:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		Any=None       ,			**_SCREAMING_SNAKE_CASE:		Tuple)				->       List[str]:

				"""simple docstring"""

				__lowerCAmelCase      :  Tuple												=       VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  List[str]												=       FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Optional[int]												=       model(input_ids=_SCREAMING_SNAKE_CASE       ,			pixel_values=_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE)

				self.assertEqual(output["text_embeds"].shape       ,			(input_ids.shape[0], config.projection_dim))
				self.assertEqual(output["image_embeds"].shape       ,			(pixel_values.shape[0], config.projection_dim))

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		str       ,			_SCREAMING_SNAKE_CASE:		Dict       ,			_SCREAMING_SNAKE_CASE:		Dict       ,			_SCREAMING_SNAKE_CASE:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]=None       ,			**_SCREAMING_SNAKE_CASE:		List[str])				->       Dict:

				"""simple docstring"""

				__lowerCAmelCase						,	__lowerCAmelCase      :  List[Any]												=       self.get_vision_text_model(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Union[str, Any]												=       {"vision_model": vision_model, "text_model": text_model}
				__lowerCAmelCase      :  Optional[Any]												=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Union[str, Any]												=       model(input_ids=_SCREAMING_SNAKE_CASE       ,			pixel_values=_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE)

				self.assertEqual(output["text_embeds"].shape       ,			(input_ids.shape[0], model.config.projection_dim))
				self.assertEqual(output["image_embeds"].shape       ,			(pixel_values.shape[0], model.config.projection_dim))

	def 							_SCREAMING_SNAKE_CASE   (     self:		Dict       ,			_SCREAMING_SNAKE_CASE:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		List[str]       ,			_SCREAMING_SNAKE_CASE:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		List[Any]=None       ,			**_SCREAMING_SNAKE_CASE:		Tuple)				->       str:

				"""simple docstring"""

				__lowerCAmelCase						,	__lowerCAmelCase      :  Optional[Any]												=       self.get_vision_text_model(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Optional[Any]												=       {"vision_model": vision_model, "text_model": text_model}
				__lowerCAmelCase      :  Any												=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Optional[Any]												=       model(input_ids=_SCREAMING_SNAKE_CASE       ,			pixel_values=_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Any												=       output[0]

				with tempfile.TemporaryDirectory() as tmpdirname:
							model.save_pretrained(_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  Union[str, Any]												=       FlaxVisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase      :  int												=       model(input_ids=_SCREAMING_SNAKE_CASE       ,			pixel_values=_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  Union[str, Any]												=       after_output[0]
							__lowerCAmelCase      :  Any												=       np.amax(np.abs(out_a - out_a))
							self.assertLessEqual(_SCREAMING_SNAKE_CASE       ,			1e-3)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Tuple       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		int       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		Any=None       ,			**_SCREAMING_SNAKE_CASE:		List[str])				->       Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase						,	__lowerCAmelCase      :  Any												=       self.get_vision_text_model(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Union[str, Any]												=       {"vision_model": vision_model, "text_model": text_model}
				__lowerCAmelCase      :  Optional[Any]												=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Optional[Any]												=       model(
				    input_ids=_SCREAMING_SNAKE_CASE       ,			pixel_values=_SCREAMING_SNAKE_CASE       ,			attention_mask=_SCREAMING_SNAKE_CASE       ,			output_attentions=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Optional[Any]												=       output.vision_model_output.attentions
				self.assertEqual(len(_SCREAMING_SNAKE_CASE)       ,			vision_config.num_hidden_layers)

				# in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
				__lowerCAmelCase      :  List[str]												=       to_atuple(vision_model.config.image_size)
				__lowerCAmelCase      :  Any												=       to_atuple(vision_model.config.patch_size)
				__lowerCAmelCase      :  Union[str, Any]												=       (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
				__lowerCAmelCase      :  Tuple												=       num_patches + 1
				self.assertEqual(vision_attentions[0].shape[-3:]       ,			(vision_config.num_attention_heads, seq_len, seq_len))

				__lowerCAmelCase      :  Union[str, Any]												=       output.text_model_output.attentions
				self.assertEqual(len(_SCREAMING_SNAKE_CASE)       ,			text_config.num_hidden_layers)

				self.assertEqual(
				    text_attentions[0].shape[-3:]       ,			(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])       ,			)

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str]       ,			_SCREAMING_SNAKE_CASE:		Union[str, Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		int)				->       str:

				"""simple docstring"""

				pt_model.to(_SCREAMING_SNAKE_CASE)
				pt_model.eval()

				# prepare inputs
				__lowerCAmelCase      :  Union[str, Any]												=       inputs_dict
				__lowerCAmelCase      :  Union[str, Any]												=       {k: torch.tensor(v.tolist()) for k, v in flax_inputs.items()}

				with torch.no_grad():
							__lowerCAmelCase      :  Any												=       pt_model(**_SCREAMING_SNAKE_CASE).to_tuple()

				__lowerCAmelCase      :  List[Any]												=       fx_model(**_SCREAMING_SNAKE_CASE).to_tuple()
				self.assertEqual(len(_SCREAMING_SNAKE_CASE)       ,			len(_SCREAMING_SNAKE_CASE)       ,			"Output lengths differ between Flax and PyTorch")
				for fx_output, pt_output in zip(fx_outputs[:4]       ,			pt_outputs[:4]):
							self.assert_almost_equals(_SCREAMING_SNAKE_CASE       ,			pt_output.numpy()       ,			4e-2)

				# PT -> Flax
				with tempfile.TemporaryDirectory() as tmpdirname:
							pt_model.save_pretrained(_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  str												=       FlaxVisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE       ,			from_pt=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  int												=       fx_model_loaded(**_SCREAMING_SNAKE_CASE).to_tuple()
				self.assertEqual(len(_SCREAMING_SNAKE_CASE)       ,			len(_SCREAMING_SNAKE_CASE)       ,			"Output lengths differ between Flax and PyTorch")
				for fx_output_loaded, pt_output in zip(fx_outputs_loaded[:4]       ,			pt_outputs[:4]):
							self.assert_almost_equals(_SCREAMING_SNAKE_CASE       ,			pt_output.numpy()       ,			4e-2)

				# Flax -> PT
				with tempfile.TemporaryDirectory() as tmpdirname:
							fx_model.save_pretrained(_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  Union[str, Any]												=       VisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE       ,			from_flax=_SCREAMING_SNAKE_CASE)

				pt_model_loaded.to(_SCREAMING_SNAKE_CASE)
				pt_model_loaded.eval()

				with torch.no_grad():
							__lowerCAmelCase      :  Optional[Any]												=       pt_model_loaded(**_SCREAMING_SNAKE_CASE).to_tuple()

				self.assertEqual(len(_SCREAMING_SNAKE_CASE)       ,			len(_SCREAMING_SNAKE_CASE)       ,			"Output lengths differ between Flax and PyTorch")
				for fx_output, pt_output_loaded in zip(fx_outputs[:4]       ,			pt_outputs_loaded[:4]):
							self.assert_almost_equals(_SCREAMING_SNAKE_CASE       ,			pt_output_loaded.numpy()       ,			4e-2)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		Dict)				->       Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  Union[str, Any]												=       VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Dict												=       VisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  str												=       FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  str												=       convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,			_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Optional[Any]												=       fx_state

				self.check_pt_flax_equivalence(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any]       ,			_SCREAMING_SNAKE_CASE:		Optional[int]       ,			_SCREAMING_SNAKE_CASE:		Tuple       ,			_SCREAMING_SNAKE_CASE:		Dict)				->       str:

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[int]												=       VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  Dict												=       VisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Tuple												=       FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase      :  int												=       load_flax_weights_in_pytorch_model(_SCREAMING_SNAKE_CASE       ,			fx_model.params)

				self.check_pt_flax_equivalence(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any])				->       str:

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       self.prepare_config_and_inputs()
				self.check_model_from_pretrained_configs(**_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		Dict)				->       int:

				"""simple docstring"""

				__lowerCAmelCase      :  List[Any]												=       self.prepare_config_and_inputs()
				self.check_vision_text_dual_encoder_from_pretrained(**_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		List[Any])				->       List[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  Dict												=       self.prepare_config_and_inputs()
				self.check_save_load(**_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE   (     self:		int)				->       Dict:

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       self.prepare_config_and_inputs()
				self.check_vision_text_output_attention(**_SCREAMING_SNAKE_CASE)

	@is_pt_flax_cross_test
	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any])				->       Any:

				"""simple docstring"""

				__lowerCAmelCase      :  Dict												=       self.prepare_config_and_inputs()
				__lowerCAmelCase      :  List[Any]												=       config_inputs_dict.pop("vision_config")
				__lowerCAmelCase      :  str												=       config_inputs_dict.pop("text_config")

				__lowerCAmelCase      :  Union[str, Any]												=       config_inputs_dict

				self.check_equivalence_pt_to_flax(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)
				self.check_equivalence_flax_to_pt(_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE       ,			_SCREAMING_SNAKE_CASE)







	@slow
	def 							_SCREAMING_SNAKE_CASE   (     self:		str)				->       Dict:

				"""simple docstring"""

				__lowerCAmelCase						,	__lowerCAmelCase      :  Dict												=       self.get_pretrained_model_and_inputs()

				__lowerCAmelCase      :  Union[str, Any]												=       model_a(**_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  List[str]												=       outputs[0]

				with tempfile.TemporaryDirectory() as tmp_dirname:
							model_a.save_pretrained(_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  List[Any]												=       FlaxVisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase      :  Union[str, Any]												=       model_a(**_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase      :  List[str]												=       after_outputs[0]
							__lowerCAmelCase      :  List[Any]												=       np.amax(np.abs(out_a - out_a))
							self.assertLessEqual(_SCREAMING_SNAKE_CASE       ,			1e-5)






@require_flax
class 			A__   (  __SCREAMING_SNAKE_CASE  ,		unittest.TestCase							):




	'''simple docstring'''





	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str])				->       List[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  Any												=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
				    "hf-internal-testing/tiny-random-vit"       ,			"hf-internal-testing/tiny-bert"       ,			vision_from_pt=_SCREAMING_SNAKE_CASE       ,			text_from_pt=_SCREAMING_SNAKE_CASE       ,			)
				__lowerCAmelCase      :  Union[str, Any]												=       13
				__lowerCAmelCase      :  Optional[int]												=       floats_tensor(
				    [
				        batch_size,
				        model.config.vision_config.num_channels,
				        model.config.vision_config.image_size,
				        model.config.vision_config.image_size,
				    ])
				__lowerCAmelCase      :  List[Any]												=       ids_tensor([batch_size, 4]       ,			model.config.text_config.vocab_size)
				__lowerCAmelCase      :  List[Any]												=       random_attention_mask([batch_size, 4])
				__lowerCAmelCase      :  str												=       {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

				return model, inputs

	def 							_SCREAMING_SNAKE_CASE   (     self:		int       ,			_SCREAMING_SNAKE_CASE:		Union[str, Any]       ,			_SCREAMING_SNAKE_CASE:		Union[str, Any])				->       Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase      :  List[str]												=       FlaxViTModel(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  Union[str, Any]												=       FlaxBertModel(_SCREAMING_SNAKE_CASE)
				return vision_model, text_model







	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[int])				->       int:

				"""simple docstring"""

				__lowerCAmelCase      :  List[Any]												=       FlaxViTModelTester(self)
				__lowerCAmelCase      :  Optional[Any]												=       FlaxBertModelTester(self)
				__lowerCAmelCase      :  int												=       vit_model_tester.prepare_config_and_inputs()
				__lowerCAmelCase      :  List[str]												=       bert_model_tester.prepare_config_and_inputs()

				__lowerCAmelCase						,	__lowerCAmelCase      :  Tuple												=       vision_config_and_inputs

				__lowerCAmelCase						,	__lowerCAmelCase						,	__lowerCAmelCase						,	__lowerCAmelCase      :  Union[str, Any]												=       text_config_and_inputs

				# make sure that cross attention layers are added
				return {
				    "text_config": text_config,
				    "vision_config": vision_config,
				    "pixel_values": pixel_values,
				    "attention_mask": attention_mask,
				    "input_ids": input_ids,
				    "token_type_ids": token_type_ids,
				}






@require_torch
class 			A__   (  __SCREAMING_SNAKE_CASE  ,		unittest.TestCase							):




	'''simple docstring'''





	def 							_SCREAMING_SNAKE_CASE   (     self:		Dict)				->       List[str]:

				"""simple docstring"""

				__lowerCAmelCase      :  Optional[int]												=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
				    "hf-internal-testing/tiny-random-clip"       ,			"hf-internal-testing/tiny-bert"       ,			vision_from_pt=_SCREAMING_SNAKE_CASE       ,			text_from_pt=_SCREAMING_SNAKE_CASE       ,			)
				__lowerCAmelCase      :  Optional[int]												=       13
				__lowerCAmelCase      :  List[str]												=       floats_tensor(
				    [
				        batch_size,
				        model.config.vision_config.num_channels,
				        model.config.vision_config.image_size,
				        model.config.vision_config.image_size,
				    ])
				__lowerCAmelCase      :  Any												=       ids_tensor([batch_size, 4]       ,			model.config.text_config.vocab_size)
				__lowerCAmelCase      :  str												=       random_attention_mask([batch_size, 4])
				__lowerCAmelCase      :  Optional[int]												=       {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

				return model, inputs

	def 							_SCREAMING_SNAKE_CASE   (     self:		Optional[Any]       ,			_SCREAMING_SNAKE_CASE:		int       ,			_SCREAMING_SNAKE_CASE:		Union[str, Any])				->       int:

				"""simple docstring"""

				__lowerCAmelCase      :  int												=       FlaxCLIPVisionModel(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase      :  List[str]												=       FlaxBertModel(_SCREAMING_SNAKE_CASE)
				return vision_model, text_model







	def 							_SCREAMING_SNAKE_CASE   (     self:		Union[str, Any])				->       Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase      :  List[Any]												=       FlaxCLIPVisionModelTester(self)
				__lowerCAmelCase      :  str												=       FlaxBertModelTester(self)
				__lowerCAmelCase      :  Optional[Any]												=       clip_model_tester.prepare_config_and_inputs()
				__lowerCAmelCase      :  Dict												=       bert_model_tester.prepare_config_and_inputs()

				__lowerCAmelCase						,	__lowerCAmelCase      :  Any												=       vision_config_and_inputs

				__lowerCAmelCase						,	__lowerCAmelCase						,	__lowerCAmelCase						,	__lowerCAmelCase      :  List[Any]												=       text_config_and_inputs

				# make sure that cross attention layers are added
				return {
				    "text_config": text_config,
				    "vision_config": vision_config,
				    "pixel_values": pixel_values,
				    "attention_mask": attention_mask,
				    "input_ids": input_ids,
				    "token_type_ids": token_type_ids,
				}






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




	'''simple docstring'''





	@slow
	def 							_SCREAMING_SNAKE_CASE   (     self:		List[str])				->       List[str]:

				"""simple docstring"""

				__lowerCAmelCase      :  Dict												=       FlaxVisionTextDualEncoderModel.from_pretrained("clip-italian/clip-italian"       ,			logit_scale_init_value=1.0)
				__lowerCAmelCase      :  str												=       VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian")

				__lowerCAmelCase      :  Tuple												=       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
				__lowerCAmelCase      :  Optional[int]												=       processor(
				    text=["una foto di un gatto", "una foto di un cane"]       ,			images=_SCREAMING_SNAKE_CASE       ,			padding=_SCREAMING_SNAKE_CASE       ,			return_tensors="np")

				__lowerCAmelCase      :  List[str]												=       model(**_SCREAMING_SNAKE_CASE)

				# verify the logits
				self.assertEqual(outputs.logits_per_image.shape       ,			(inputs.pixel_values.shape[0], inputs.input_ids.shape[0]))
				self.assertEqual(
				    outputs.logits_per_text.shape       ,			(inputs.input_ids.shape[0], inputs.pixel_values.shape[0])       ,			)

				__lowerCAmelCase      :  List[str]												=       np.array([[1.228_4727, 0.310_4122]])

				self.assertTrue(np.allclose(outputs.logits_per_image       ,			_SCREAMING_SNAKE_CASE       ,			atol=1e-3))