


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




					'''simple docstring'''


					if isinstance(lowerCAmelCase_  ,						collections.abc.Iterable   ):
										return x
					return (x, x)




@require_flax
class  UpperCamelCase_	:

	"""simple docstring"""





	def   UpperCAmelCase_      ( self			:				List[Any]			,			UpperCAmelCase__			:				Union[str, Any]			,			UpperCAmelCase__			:				Tuple   )       ->	Tuple:
						pass




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




	def   UpperCAmelCase_      ( self			:				List[str]   )       ->	Optional[Any]:
						pass




	def   UpperCAmelCase_      ( self			:				Tuple			,			UpperCAmelCase__			:				np.ndarray			,			UpperCAmelCase__			:				np.ndarray			,			UpperCAmelCase__			:				float   )       ->	List[Any]:
						__SCREAMING_SNAKE_CASE						   =				np.abs((a - b)   ).max()
						self.assertLessEqual(UpperCAmelCase__			,			UpperCAmelCase__			,			F"""Difference between torch and flax is {diff} (>= {tol})."""   )




	def   UpperCAmelCase_      ( self			:				Any			,			UpperCAmelCase__			:				Any			,			UpperCAmelCase__			:				str			,			UpperCAmelCase__			:				str			,			UpperCAmelCase__			:				Optional[int]			,			UpperCAmelCase__			:				Any=None			,			**UpperCAmelCase__			:				Union[str, Any]   )       ->	Dict:
						__SCREAMING_SNAKE_CASE						   =				VisionTextDualEncoderConfig.from_vision_text_configs(UpperCAmelCase__			,			UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				FlaxVisionTextDualEncoderModel(UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				model(input_ids=UpperCAmelCase__			,			pixel_values=UpperCAmelCase__			,			attention_mask=UpperCAmelCase__   )

						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   UpperCAmelCase_      ( self			:				List[str]			,			UpperCAmelCase__			:				List[Any]			,			UpperCAmelCase__			:				Union[str, Any]			,			UpperCAmelCase__			:				List[str]			,			UpperCAmelCase__			:				Optional[Any]			,			UpperCAmelCase__			:				Dict=None			,			**UpperCAmelCase__			:				str   )       ->	Tuple:
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				self.get_vision_text_model(UpperCAmelCase__			,			UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				{"vision_model": vision_model, "text_model": text_model}
						__SCREAMING_SNAKE_CASE						   =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				model(input_ids=UpperCAmelCase__			,			pixel_values=UpperCAmelCase__			,			attention_mask=UpperCAmelCase__   )

						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   UpperCAmelCase_      ( self			:				Optional[Any]			,			UpperCAmelCase__			:				str			,			UpperCAmelCase__			:				int			,			UpperCAmelCase__			:				int			,			UpperCAmelCase__			:				Tuple			,			UpperCAmelCase__			:				int=None			,			**UpperCAmelCase__			:				Dict   )       ->	Tuple:
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				self.get_vision_text_model(UpperCAmelCase__			,			UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				{"vision_model": vision_model, "text_model": text_model}
						__SCREAMING_SNAKE_CASE						   =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				model(input_ids=UpperCAmelCase__			,			pixel_values=UpperCAmelCase__			,			attention_mask=UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				output[0]

						with tempfile.TemporaryDirectory() as tmpdirname:
											model.save_pretrained(UpperCAmelCase__   )
											__SCREAMING_SNAKE_CASE						   =				FlaxVisionTextDualEncoderModel.from_pretrained(UpperCAmelCase__   )

											__SCREAMING_SNAKE_CASE						   =				model(input_ids=UpperCAmelCase__			,			pixel_values=UpperCAmelCase__			,			attention_mask=UpperCAmelCase__   )
											__SCREAMING_SNAKE_CASE						   =				after_output[0]
											__SCREAMING_SNAKE_CASE						   =				np.amax(np.abs(out_a - out_a   )   )
											self.assertLessEqual(UpperCAmelCase__			,			1E-3   )




	def   UpperCAmelCase_      ( self			:				List[Any]			,			UpperCAmelCase__			:				Optional[Any]			,			UpperCAmelCase__			:				int			,			UpperCAmelCase__			:				Optional[int]			,			UpperCAmelCase__			:				Any			,			UpperCAmelCase__			:				Tuple=None			,			**UpperCAmelCase__			:				Any   )       ->	Union[str, Any]:
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				self.get_vision_text_model(UpperCAmelCase__			,			UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				{"vision_model": vision_model, "text_model": text_model}
						__SCREAMING_SNAKE_CASE						   =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				model(
						    input_ids=UpperCAmelCase__			,			pixel_values=UpperCAmelCase__			,			attention_mask=UpperCAmelCase__			,			output_attentions=UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				output.vision_model_output.attentions
						self.assertEqual(len(UpperCAmelCase__   )			,			vision_config.num_hidden_layers   )

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

						__SCREAMING_SNAKE_CASE						   =				output.text_model_output.attentions
						self.assertEqual(len(UpperCAmelCase__   )			,			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   UpperCAmelCase_      ( self			:				Any			,			UpperCAmelCase__			:				Dict			,			UpperCAmelCase__			:				List[Any]			,			UpperCAmelCase__			:				List[str]   )       ->	int:
						pt_model.to(UpperCAmelCase__   )
						pt_model.eval()

						# prepare inputs
						__SCREAMING_SNAKE_CASE						   =				inputs_dict
						__SCREAMING_SNAKE_CASE						   =				{k: torch.tensor(v.tolist()   ) for k, v in flax_inputs.items()}

						with torch.no_grad():
											__SCREAMING_SNAKE_CASE						   =				pt_model(**UpperCAmelCase__   ).to_tuple()

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

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

						__SCREAMING_SNAKE_CASE						   =				fx_model_loaded(**UpperCAmelCase__   ).to_tuple()
						self.assertEqual(len(UpperCAmelCase__   )			,			len(UpperCAmelCase__   )			,			"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(UpperCAmelCase__			,			pt_output.numpy()			,			4E-2   )

						# Flax -> PT
						with tempfile.TemporaryDirectory() as tmpdirname:
											fx_model.save_pretrained(UpperCAmelCase__   )
											__SCREAMING_SNAKE_CASE						   =				VisionTextDualEncoderModel.from_pretrained(UpperCAmelCase__			,			from_flax=UpperCAmelCase__   )

						pt_model_loaded.to(UpperCAmelCase__   )
						pt_model_loaded.eval()

						with torch.no_grad():
											__SCREAMING_SNAKE_CASE						   =				pt_model_loaded(**UpperCAmelCase__   ).to_tuple()

						self.assertEqual(len(UpperCAmelCase__   )			,			len(UpperCAmelCase__   )			,			"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(UpperCAmelCase__			,			pt_output_loaded.numpy()			,			4E-2   )




	def   UpperCAmelCase_      ( self			:				Dict			,			UpperCAmelCase__			:				Tuple			,			UpperCAmelCase__			:				int			,			UpperCAmelCase__			:				Any   )       ->	int:
						__SCREAMING_SNAKE_CASE						   =				VisionTextDualEncoderConfig.from_vision_text_configs(UpperCAmelCase__			,			UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				VisionTextDualEncoderModel(UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				FlaxVisionTextDualEncoderModel(UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				convert_pytorch_state_dict_to_flax(pt_model.state_dict()			,			UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				fx_state

						self.check_pt_flax_equivalence(UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__   )




	def   UpperCAmelCase_      ( self			:				Union[str, Any]			,			UpperCAmelCase__			:				List[str]			,			UpperCAmelCase__			:				Dict			,			UpperCAmelCase__			:				str   )       ->	Optional[int]:
						__SCREAMING_SNAKE_CASE						   =				VisionTextDualEncoderConfig.from_vision_text_configs(UpperCAmelCase__			,			UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				VisionTextDualEncoderModel(UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				FlaxVisionTextDualEncoderModel(UpperCAmelCase__   )

						__SCREAMING_SNAKE_CASE						   =				load_flax_weights_in_pytorch_model(UpperCAmelCase__			,			fx_model.params   )

						self.check_pt_flax_equivalence(UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__   )




	def   UpperCAmelCase_      ( self			:				Dict   )       ->	Tuple:
						__SCREAMING_SNAKE_CASE						   =				self.prepare_config_and_inputs()
						self.check_model_from_pretrained_configs(**UpperCAmelCase__   )




	def   UpperCAmelCase_      ( self			:				Any   )       ->	Tuple:
						__SCREAMING_SNAKE_CASE						   =				self.prepare_config_and_inputs()
						self.check_vision_text_dual_encoder_from_pretrained(**UpperCAmelCase__   )




	def   UpperCAmelCase_      ( self			:				int   )       ->	Tuple:
						__SCREAMING_SNAKE_CASE						   =				self.prepare_config_and_inputs()
						self.check_save_load(**UpperCAmelCase__   )




	def   UpperCAmelCase_      ( self			:				Optional[Any]   )       ->	int:
						__SCREAMING_SNAKE_CASE						   =				self.prepare_config_and_inputs()
						self.check_vision_text_output_attention(**UpperCAmelCase__   )




	@is_pt_flax_cross_test
	def   UpperCAmelCase_      ( self			:				Optional[int]   )       ->	List[Any]:
						__SCREAMING_SNAKE_CASE						   =				self.prepare_config_and_inputs()
						__SCREAMING_SNAKE_CASE						   =				config_inputs_dict.pop("vision_config"   )
						__SCREAMING_SNAKE_CASE						   =				config_inputs_dict.pop("text_config"   )

						__SCREAMING_SNAKE_CASE						   =				config_inputs_dict

						self.check_equivalence_pt_to_flax(UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__   )
						self.check_equivalence_flax_to_pt(UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__   )




	@slow
	def   UpperCAmelCase_      ( self			:				Optional[Any]   )       ->	Union[str, Any]:
						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				self.get_pretrained_model_and_inputs()

						__SCREAMING_SNAKE_CASE						   =				model_a(**UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				outputs[0]

						with tempfile.TemporaryDirectory() as tmp_dirname:
											model_a.save_pretrained(UpperCAmelCase__   )
											__SCREAMING_SNAKE_CASE						   =				FlaxVisionTextDualEncoderModel.from_pretrained(UpperCAmelCase__   )

											__SCREAMING_SNAKE_CASE						   =				model_a(**UpperCAmelCase__   )
											__SCREAMING_SNAKE_CASE						   =				after_outputs[0]
											__SCREAMING_SNAKE_CASE						   =				np.amax(np.abs(out_a - out_a   )   )
											self.assertLessEqual(UpperCAmelCase__			,			1E-5   )




@require_flax
class  UpperCamelCase_	( UpperCamelCase			,  unittest.TestCase):

	"""simple docstring"""





	def   UpperCAmelCase_      ( self			:				Tuple   )       ->	int:
						__SCREAMING_SNAKE_CASE						   =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
						    "hf-internal-testing/tiny-random-vit"			,			"hf-internal-testing/tiny-bert"			,			vision_from_pt=UpperCAmelCase__			,			text_from_pt=UpperCAmelCase__			,			)
						__SCREAMING_SNAKE_CASE						   =				1_3
						__SCREAMING_SNAKE_CASE						   =				floats_tensor(
						    [
						        batch_size,
						        model.config.vision_config.num_channels,
						        model.config.vision_config.image_size,
						        model.config.vision_config.image_size,
						    ]   )
						__SCREAMING_SNAKE_CASE						   =				ids_tensor([batch_size, 4]			,			model.config.text_config.vocab_size   )
						__SCREAMING_SNAKE_CASE						   =				random_attention_mask([batch_size, 4]   )
						__SCREAMING_SNAKE_CASE						   =				{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

						return model, inputs




	def   UpperCAmelCase_      ( self			:				List[str]			,			UpperCAmelCase__			:				int			,			UpperCAmelCase__			:				List[str]   )       ->	Dict:
						__SCREAMING_SNAKE_CASE						   =				FlaxViTModel(UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				FlaxBertModel(UpperCAmelCase__   )
						return vision_model, text_model




	def   UpperCAmelCase_      ( self			:				Dict   )       ->	Dict:
						__SCREAMING_SNAKE_CASE						   =				FlaxViTModelTester(self   )
						__SCREAMING_SNAKE_CASE						   =				FlaxBertModelTester(self   )
						__SCREAMING_SNAKE_CASE						   =				vit_model_tester.prepare_config_and_inputs()
						__SCREAMING_SNAKE_CASE						   =				bert_model_tester.prepare_config_and_inputs()

						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				vision_config_and_inputs

						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				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  UpperCamelCase_	( UpperCamelCase			,  unittest.TestCase):

	"""simple docstring"""





	def   UpperCAmelCase_      ( self			:				List[str]   )       ->	List[Any]:
						__SCREAMING_SNAKE_CASE						   =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
						    "hf-internal-testing/tiny-random-clip"			,			"hf-internal-testing/tiny-bert"			,			vision_from_pt=UpperCAmelCase__			,			text_from_pt=UpperCAmelCase__			,			)
						__SCREAMING_SNAKE_CASE						   =				1_3
						__SCREAMING_SNAKE_CASE						   =				floats_tensor(
						    [
						        batch_size,
						        model.config.vision_config.num_channels,
						        model.config.vision_config.image_size,
						        model.config.vision_config.image_size,
						    ]   )
						__SCREAMING_SNAKE_CASE						   =				ids_tensor([batch_size, 4]			,			model.config.text_config.vocab_size   )
						__SCREAMING_SNAKE_CASE						   =				random_attention_mask([batch_size, 4]   )
						__SCREAMING_SNAKE_CASE						   =				{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

						return model, inputs




	def   UpperCAmelCase_      ( self			:				Dict			,			UpperCAmelCase__			:				List[str]			,			UpperCAmelCase__			:				List[Any]   )       ->	Dict:
						__SCREAMING_SNAKE_CASE						   =				FlaxCLIPVisionModel(UpperCAmelCase__   )
						__SCREAMING_SNAKE_CASE						   =				FlaxBertModel(UpperCAmelCase__   )
						return vision_model, text_model




	def   UpperCAmelCase_      ( self			:				Optional[Any]   )       ->	Tuple:
						__SCREAMING_SNAKE_CASE						   =				FlaxCLIPVisionModelTester(self   )
						__SCREAMING_SNAKE_CASE						   =				FlaxBertModelTester(self   )
						__SCREAMING_SNAKE_CASE						   =				clip_model_tester.prepare_config_and_inputs()
						__SCREAMING_SNAKE_CASE						   =				bert_model_tester.prepare_config_and_inputs()

						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				vision_config_and_inputs

						__SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE   ,  __SCREAMING_SNAKE_CASE						   =				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  UpperCamelCase_	( unittest.TestCase):

	"""simple docstring"""





	@slow
	def   UpperCAmelCase_      ( self			:				Any   )       ->	List[str]:
						__SCREAMING_SNAKE_CASE						   =				FlaxVisionTextDualEncoderModel.from_pretrained("clip-italian/clip-italian"			,			logit_scale_init_value=1.0   )
						__SCREAMING_SNAKE_CASE						   =				VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian"   )

						__SCREAMING_SNAKE_CASE						   =				Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"   )
						__SCREAMING_SNAKE_CASE						   =				processor(
						    text=["una foto di un gatto", "una foto di un cane"]			,			images=UpperCAmelCase__			,			padding=UpperCAmelCase__			,			return_tensors="np"   )

						__SCREAMING_SNAKE_CASE						   =				model(**UpperCAmelCase__   )

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

						__SCREAMING_SNAKE_CASE						   =				np.array([[1.2_284_727, 0.3_104_122]]   )

						self.assertTrue(np.allclose(outputs.logits_per_image			,			UpperCAmelCase__			,			atol=1E-3   )   )




"""simple docstring"""







from collections import defaultdict
from math import ceil, sqrt




def    UpperCAmelCase__       (lowerCAmelCase_ = 100_0000  ,						lowerCAmelCase_ = 10   ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						   =				defaultdict(lowerCAmelCase_   )

					for outer_width in range(3  ,						(t_limit // 4) + 2   ):
										if outer_width * outer_width > t_limit:
															__SCREAMING_SNAKE_CASE						   =				max(
															    ceil(sqrt(outer_width * outer_width - t_limit   )   )  ,						1   )
										else:
															__SCREAMING_SNAKE_CASE						   =				1

										hole_width_lower_bound += (outer_width - hole_width_lower_bound) % 2

										for hole_width in range(lowerCAmelCase_  ,						outer_width - 1  ,						2   ):
															count[outer_width * outer_width - hole_width * hole_width] += 1

					return sum(1 for n in count.values() if 1 <= n <= 10   )


if __name__ == "__main__":
		print(F"{solution() = }")
