
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 				lowercase		:




							def  A__      (				self   ,A__   ,A__):
												pass




							def  A__      (				self):
												pass




							def  A__      (				self):
												pass




							def  A__      (				self   ,A__   ,A__   ,A__):
												lowercase							=						np.abs((a - b)).max()
												self.assertLessEqual(A__   ,A__   ,f'Difference between torch and flax is {diff} (>= {tol}).')




							def  A__      (				self   ,A__   ,A__   ,A__   ,A__   ,A__=None   ,**A__):
												lowercase							=						VisionTextDualEncoderConfig.from_vision_text_configs(A__   ,A__)

												lowercase							=						FlaxVisionTextDualEncoderModel(A__)

												lowercase							=						model(input_ids=A__   ,pixel_values=A__   ,attention_mask=A__)

												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  A__      (				self   ,A__   ,A__   ,A__   ,A__   ,A__=None   ,**A__):
												lowercase	,						lowercase							=						self.get_vision_text_model(A__   ,A__)
												lowercase							=						{'''vision_model''': vision_model, '''text_model''': text_model}
												lowercase							=						FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**A__)

												lowercase							=						model(input_ids=A__   ,pixel_values=A__   ,attention_mask=A__)

												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  A__      (				self   ,A__   ,A__   ,A__   ,A__   ,A__=None   ,**A__):
												lowercase	,						lowercase							=						self.get_vision_text_model(A__   ,A__)
												lowercase							=						{'''vision_model''': vision_model, '''text_model''': text_model}
												lowercase							=						FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**A__)

												lowercase							=						model(input_ids=A__   ,pixel_values=A__   ,attention_mask=A__)
												lowercase							=						output[0]

												with tempfile.TemporaryDirectory() as tmpdirname:
																	model.save_pretrained(A__)
																	lowercase							=						FlaxVisionTextDualEncoderModel.from_pretrained(A__)

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




							def  A__      (				self   ,A__   ,A__   ,A__   ,A__   ,A__=None   ,**A__):
												lowercase	,						lowercase							=						self.get_vision_text_model(A__   ,A__)
												lowercase							=						{'''vision_model''': vision_model, '''text_model''': text_model}
												lowercase							=						FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**A__)

												lowercase							=						model(
												    input_ids=A__   ,pixel_values=A__   ,attention_mask=A__   ,output_attentions=A__)

												lowercase							=						output.vision_model_output.attentions
												self.assertEqual(len(A__)   ,vision_config.num_hidden_layers)

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

												lowercase							=						output.text_model_output.attentions
												self.assertEqual(len(A__)   ,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  A__      (				self   ,A__   ,A__   ,A__):
												pt_model.to(A__)
												pt_model.eval()

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

												with torch.no_grad():
																	lowercase							=						pt_model(**A__).to_tuple()

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

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

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

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

												pt_model_loaded.to(A__)
												pt_model_loaded.eval()

												with torch.no_grad():
																	lowercase							=						pt_model_loaded(**A__).to_tuple()

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




							def  A__      (				self   ,A__   ,A__   ,A__):
												lowercase							=						VisionTextDualEncoderConfig.from_vision_text_configs(A__   ,A__)

												lowercase							=						VisionTextDualEncoderModel(A__)
												lowercase							=						FlaxVisionTextDualEncoderModel(A__)

												lowercase							=						convert_pytorch_state_dict_to_flax(pt_model.state_dict()   ,A__)
												lowercase							=						fx_state

												self.check_pt_flax_equivalence(A__   ,A__   ,A__)




							def  A__      (				self   ,A__   ,A__   ,A__):
												lowercase							=						VisionTextDualEncoderConfig.from_vision_text_configs(A__   ,A__)

												lowercase							=						VisionTextDualEncoderModel(A__)
												lowercase							=						FlaxVisionTextDualEncoderModel(A__)

												lowercase							=						load_flax_weights_in_pytorch_model(A__   ,fx_model.params)

												self.check_pt_flax_equivalence(A__   ,A__   ,A__)




							def  A__      (				self):
												lowercase							=						self.prepare_config_and_inputs()
												self.check_model_from_pretrained_configs(**A__)




							def  A__      (				self):
												lowercase							=						self.prepare_config_and_inputs()
												self.check_vision_text_dual_encoder_from_pretrained(**A__)




							def  A__      (				self):
												lowercase							=						self.prepare_config_and_inputs()
												self.check_save_load(**A__)




							def  A__      (				self):
												lowercase							=						self.prepare_config_and_inputs()
												self.check_vision_text_output_attention(**A__)




							@is_pt_flax_cross_test
							def  A__      (				self):
												lowercase							=						self.prepare_config_and_inputs()
												lowercase							=						config_inputs_dict.pop('''vision_config''')
												lowercase							=						config_inputs_dict.pop('''text_config''')

												lowercase							=						config_inputs_dict

												self.check_equivalence_pt_to_flax(A__   ,A__   ,A__)
												self.check_equivalence_flax_to_pt(A__   ,A__   ,A__)





							@slow
							def  A__      (				self):
												lowercase	,						lowercase							=						self.get_pretrained_model_and_inputs()

												lowercase							=						model_a(**A__)
												lowercase							=						outputs[0]

												with tempfile.TemporaryDirectory() as tmp_dirname:
																	model_a.save_pretrained(A__)
																	lowercase							=						FlaxVisionTextDualEncoderModel.from_pretrained(A__)

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







@require_flax
class 				lowercase		(					SCREAMING_SNAKE_CASE__       ,  unittest.TestCase		):




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

												return model, inputs




							def  A__      (				self   ,A__   ,A__):
												lowercase							=						FlaxViTModel(A__)
												lowercase							=						FlaxBertModel(A__)
												return vision_model, text_model





							def  A__      (				self):
												lowercase							=						FlaxViTModelTester(self)
												lowercase							=						FlaxBertModelTester(self)
												lowercase							=						vit_model_tester.prepare_config_and_inputs()
												lowercase							=						bert_model_tester.prepare_config_and_inputs()

												lowercase	,						lowercase							=						vision_config_and_inputs

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




							def  A__      (				self):
												lowercase							=						FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
												    '''hf-internal-testing/tiny-random-clip'''   ,'''hf-internal-testing/tiny-bert'''   ,vision_from_pt=A__   ,text_from_pt=A__   ,)
												lowercase							=						1_3
												lowercase							=						floats_tensor(
												    [
												        batch_size,
												        model.config.vision_config.num_channels,
												        model.config.vision_config.image_size,
												        model.config.vision_config.image_size,
												    ])
												lowercase							=						ids_tensor([batch_size, 4]   ,model.config.text_config.vocab_size)
												lowercase							=						random_attention_mask([batch_size, 4])
												lowercase							=						{'''pixel_values''': pixel_values, '''input_ids''': input_ids, '''attention_mask''': attention_mask}

												return model, inputs




							def  A__      (				self   ,A__   ,A__):
												lowercase							=						FlaxCLIPVisionModel(A__)
												lowercase							=						FlaxBertModel(A__)
												return vision_model, text_model





							def  A__      (				self):
												lowercase							=						FlaxCLIPVisionModelTester(self)
												lowercase							=						FlaxBertModelTester(self)
												lowercase							=						clip_model_tester.prepare_config_and_inputs()
												lowercase							=						bert_model_tester.prepare_config_and_inputs()

												lowercase	,						lowercase							=						vision_config_and_inputs

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




							@slow
							def  A__      (				self):
												lowercase							=						FlaxVisionTextDualEncoderModel.from_pretrained('''clip-italian/clip-italian'''   ,logit_scale_init_value=1.0)
												lowercase							=						VisionTextDualEncoderProcessor.from_pretrained('''clip-italian/clip-italian''')

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

												lowercase							=						model(**A__)

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

												lowercase							=						np.array([[1.2284727, 0.3104122]])

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





'''simple docstring'''

def 			__UpperCamelCase       (   UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase     ):

 def update_area_of_max_square(UpperCAmelCase		,     UpperCAmelCase     ) -> int:
  # BASE CASE
  if row >= rows or col >= cols:
   return 0

  lowercase__				:    int  =					update_area_of_max_square(UpperCAmelCase		,     col + 1     )
  lowercase__				:    Any  =					update_area_of_max_square(row + 1		,     col + 1     )
  lowercase__				:    str  =					update_area_of_max_square(row + 1		,     UpperCAmelCase     )

  if mat[row][col]:
   lowercase__				:    List[Any]  =					1 + min([right, diagonal, down]     )
   lowercase__				:    List[Any]  =					max(largest_square_area[0]		,     UpperCAmelCase     )
   return sub_problem_sol
  else:
   return 0

 lowercase__				:    Dict  =					[0]
 update_area_of_max_square(0		,     0     )
 return largest_square_area[0]
def 			__UpperCamelCase       (   UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase     ):

 def update_area_of_max_square_using_dp_array(
     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase     ) -> int:
  if row >= rows or col >= cols:
   return 0
  if dp_array[row][col] != -1:
   return dp_array[row][col]

  lowercase__				:    int  =					update_area_of_max_square_using_dp_array(UpperCAmelCase		,     col + 1		,     UpperCAmelCase     )
  lowercase__				:    Optional[int]  =					update_area_of_max_square_using_dp_array(row + 1		,     col + 1		,     UpperCAmelCase     )
  lowercase__				:    Any  =					update_area_of_max_square_using_dp_array(row + 1		,     UpperCAmelCase		,     UpperCAmelCase     )

  if mat[row][col]:
   lowercase__				:    Optional[int]  =					1 + min([right, diagonal, down]     )
   lowercase__				:    Any  =					max(largest_square_area[0]		,     UpperCAmelCase     )
   lowercase__				:    int  =					sub_problem_sol
   return sub_problem_sol
  else:
   return 0

 lowercase__				:    Any  =					[0]
 lowercase__				:    List[Any]  =					[[-1] * cols for _ in range(UpperCAmelCase     )]
 update_area_of_max_square_using_dp_array(0		,     0		,     UpperCAmelCase     )

 return largest_square_area[0]
def 			__UpperCamelCase       (   UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase     ):
 lowercase__				:    Optional[int]  =					[[0] * (cols + 1) for _ in range(rows + 1     )]
 lowercase__				:    str  =					0
 for row in range(rows - 1		,     -1		,     -1     ):
  for col in range(cols - 1		,     -1		,     -1     ):
   lowercase__				:    str  =					dp_array[row][col + 1]
   lowercase__				:    Optional[Any]  =					dp_array[row + 1][col + 1]
   lowercase__				:    str  =					dp_array[row + 1][col]

   if mat[row][col] == 1:
    lowercase__				:    Dict  =					1 + min(UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase     )
    lowercase__				:    str  =					max(dp_array[row][col]		,     UpperCAmelCase     )
   else:
    lowercase__				:    Any  =					0

 return largest_square_area
def 			__UpperCamelCase       (   UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase     ):
 lowercase__				:    List[str]  =					[0] * (cols + 1)
 lowercase__				:    str  =					[0] * (cols + 1)
 lowercase__				:    Tuple  =					0
 for row in range(rows - 1		,     -1		,     -1     ):
  for col in range(cols - 1		,     -1		,     -1     ):
   lowercase__				:    List[Any]  =					current_row[col + 1]
   lowercase__				:    Any  =					next_row[col + 1]
   lowercase__				:    Optional[Any]  =					next_row[col]

   if mat[row][col] == 1:
    lowercase__				:    List[str]  =					1 + min(UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase     )
    lowercase__				:    List[str]  =					max(current_row[col]		,     UpperCAmelCase     )
   else:
    lowercase__				:    int  =					0
  lowercase__				:    int  =					current_row

 return largest_square_area


if __name__ == "__main__":
  import doctest

  doctest.testmod()
  print(largest_square_area_in_matrix_bottom_up(2, 2, [[1, 1], [1, 1]]))

