






import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      MODEL_FOR_PRETRAINING_MAPPING,
      MegatronBertForCausalLM,
      MegatronBertForMaskedLM,
      MegatronBertForMultipleChoice,
      MegatronBertForNextSentencePrediction,
      MegatronBertForPreTraining,
      MegatronBertForQuestionAnswering,
      MegatronBertForSequenceClassification,
      MegatronBertForTokenClassification,
      MegatronBertModel,
  )
class 		SCREAMING_SNAKE_CASE__	:

  '''simple docstring'''







  def __init__(      self		:			List[str]       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			List[Any]=13       ,				lowerCamelCase		:			List[Any]=7       ,				lowerCamelCase		:			int=True       ,				lowerCamelCase		:			int=True       ,				lowerCamelCase		:			Dict=True       ,				lowerCamelCase		:			Optional[int]=True       ,				lowerCamelCase		:			Optional[Any]=99       ,				lowerCamelCase		:			Tuple=64       ,				lowerCamelCase		:			List[Any]=32       ,				lowerCamelCase		:			Tuple=5       ,				lowerCamelCase		:			Tuple=4       ,				lowerCamelCase		:			Union[str, Any]=37       ,				lowerCamelCase		:			List[Any]="gelu"       ,				lowerCamelCase		:			int=0.1       ,				lowerCamelCase		:			int=0.1       ,				lowerCamelCase		:			List[Any]=512       ,				lowerCamelCase		:			Tuple=16       ,				lowerCamelCase		:			Optional[Any]=2       ,				lowerCamelCase		:			int=0.02       ,				lowerCamelCase		:			Any=3       ,				lowerCamelCase		:			Dict=4       ,				lowerCamelCase		:			Tuple=None       ,				)						->	int:


         """simple docstring"""
         _UpperCAmelCase							=    parent
         _UpperCAmelCase							=    batch_size
         _UpperCAmelCase							=    seq_length
         _UpperCAmelCase							=    is_training
         _UpperCAmelCase							=    use_input_mask
         _UpperCAmelCase							=    use_token_type_ids
         _UpperCAmelCase							=    use_labels
         _UpperCAmelCase							=    vocab_size
         _UpperCAmelCase							=    hidden_size
         _UpperCAmelCase							=    embedding_size
         _UpperCAmelCase							=    num_hidden_layers
         _UpperCAmelCase							=    num_attention_heads
         _UpperCAmelCase							=    intermediate_size
         _UpperCAmelCase							=    hidden_act
         _UpperCAmelCase							=    hidden_dropout_prob
         _UpperCAmelCase							=    attention_probs_dropout_prob
         _UpperCAmelCase							=    max_position_embeddings
         _UpperCAmelCase							=    type_vocab_size
         _UpperCAmelCase							=    type_sequence_label_size
         _UpperCAmelCase							=    initializer_range
         _UpperCAmelCase							=    num_labels
         _UpperCAmelCase							=    num_choices
         _UpperCAmelCase							=    scope


  def  lowerCamelCase       (      self		:			Union[str, Any]							)						->	List[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    ids_tensor([self.batch_size, self.seq_length]       ,				self.vocab_size							)

         _UpperCAmelCase							=    None
         if self.use_input_mask:
                _UpperCAmelCase							=    random_attention_mask([self.batch_size, self.seq_length]							)

         _UpperCAmelCase							=    None
         if self.use_token_type_ids:
                _UpperCAmelCase							=    ids_tensor([self.batch_size, self.seq_length]       ,				self.type_vocab_size							)

         _UpperCAmelCase							=    None
         _UpperCAmelCase							=    None
         _UpperCAmelCase							=    None
         if self.use_labels:
                _UpperCAmelCase							=    ids_tensor([self.batch_size]       ,				self.type_sequence_label_size							)
                _UpperCAmelCase							=    ids_tensor([self.batch_size, self.seq_length]       ,				self.num_labels							)
                _UpperCAmelCase							=    ids_tensor([self.batch_size]       ,				self.num_choices							)

         _UpperCAmelCase							=    self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


  def  lowerCamelCase       (      self		:			Any							)						->	Any:


         """simple docstring"""
         return MegatronBertConfig(
             vocab_size=self.vocab_size       ,				hidden_size=self.hidden_size       ,				num_hidden_layers=self.num_hidden_layers       ,				num_attention_heads=self.num_attention_heads       ,				intermediate_size=self.intermediate_size       ,				embedding_size=self.embedding_size       ,				hidden_act=self.hidden_act       ,				hidden_dropout_prob=self.hidden_dropout_prob       ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,				max_position_embeddings=self.max_position_embeddings       ,				type_vocab_size=self.type_vocab_size       ,				is_decoder=lowerCamelCase       ,				initializer_range=self.initializer_range       ,				)


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


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertModel(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase							)
         _UpperCAmelCase							=    model(lowerCamelCase       ,				token_type_ids=lowerCamelCase							)
         _UpperCAmelCase							=    model(lowerCamelCase							)

         self.parent.assertEqual(result.last_hidden_state.shape       ,				(self.batch_size, self.seq_length, self.hidden_size)							)
         self.parent.assertEqual(result.pooler_output.shape       ,				(self.batch_size, self.hidden_size)							)


  def  lowerCamelCase       (      self		:			List[Any]       ,				lowerCamelCase		:			str       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Any       ,				lowerCamelCase		:			str       ,				lowerCamelCase		:			List[Any]							)						->	int:


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertForMaskedLM(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase       ,				labels=lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.seq_length, self.vocab_size)							)


  def  lowerCamelCase       (      self		:			Tuple       ,				lowerCamelCase		:			str       ,				lowerCamelCase		:			Tuple       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			List[str]       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			Tuple       ,				lowerCamelCase		:			Union[str, Any]							)						->	List[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertForCausalLM(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase       ,				labels=lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.seq_length, self.vocab_size)							)


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


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertForNextSentencePrediction(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(
             lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase       ,				labels=lowerCamelCase       ,				)
         self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, 2)							)


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


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertForPreTraining(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(
             lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase       ,				labels=lowerCamelCase       ,				next_sentence_label=lowerCamelCase       ,				)
         self.parent.assertEqual(result.prediction_logits.shape       ,				(self.batch_size, self.seq_length, self.vocab_size)							)
         self.parent.assertEqual(result.seq_relationship_logits.shape       ,				(self.batch_size, 2)							)


  def  lowerCamelCase       (      self		:			Any       ,				lowerCamelCase		:			List[str]       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			List[str]       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			Any       ,				lowerCamelCase		:			int							)						->	List[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertForQuestionAnswering(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(
             lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase       ,				start_positions=lowerCamelCase       ,				end_positions=lowerCamelCase       ,				)
         self.parent.assertEqual(result.start_logits.shape       ,				(self.batch_size, self.seq_length)							)
         self.parent.assertEqual(result.end_logits.shape       ,				(self.batch_size, self.seq_length)							)


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


         """simple docstring"""
         _UpperCAmelCase							=    self.num_labels
         _UpperCAmelCase							=    MegatronBertForSequenceClassification(lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase       ,				labels=lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.num_labels)							)


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


         """simple docstring"""
         _UpperCAmelCase							=    self.num_labels
         _UpperCAmelCase							=    MegatronBertForTokenClassification(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase       ,				labels=lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.seq_length, self.num_labels)							)


  def  lowerCamelCase       (      self		:			str       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			List[str]       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Tuple							)						->	List[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    self.num_choices
         _UpperCAmelCase							=    MegatronBertForMultipleChoice(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    input_ids.unsqueeze(1							).expand(-1       ,				self.num_choices       ,				-1							).contiguous()
         _UpperCAmelCase							=    token_type_ids.unsqueeze(1							).expand(-1       ,				self.num_choices       ,				-1							).contiguous()
         _UpperCAmelCase							=    input_mask.unsqueeze(1							).expand(-1       ,				self.num_choices       ,				-1							).contiguous()
         _UpperCAmelCase							=    model(
             lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase       ,				labels=lowerCamelCase       ,				)
         self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.num_choices)							)





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


         """simple docstring"""
         _UpperCAmelCase							=    self.prepare_config_and_inputs()
         (
             (
             _UpperCAmelCase
         )	,			(
             _UpperCAmelCase
         )	,			(
             _UpperCAmelCase
         )	,			(
             _UpperCAmelCase
         )	,			(
             _UpperCAmelCase
         )	,			(
             _UpperCAmelCase
         )	,			(
             _UpperCAmelCase
         )	,			
         )							=    config_and_inputs
         _UpperCAmelCase							=    {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
         return config, inputs_dict



@require_torch
class 		SCREAMING_SNAKE_CASE__	(  UpperCAmelCase					,      UpperCAmelCase					,      unittest.TestCase     ):

  '''simple docstring'''







  _lowerCamelCase    =  (
      (
          MegatronBertModel,
          MegatronBertForMaskedLM,
          MegatronBertForCausalLM,
          MegatronBertForMultipleChoice,
          MegatronBertForNextSentencePrediction,
          MegatronBertForPreTraining,
          MegatronBertForQuestionAnswering,
          MegatronBertForSequenceClassification,
          MegatronBertForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  _lowerCamelCase    =  (
      {
          '''feature-extraction''': MegatronBertModel,
          '''fill-mask''': MegatronBertForMaskedLM,
          '''question-answering''': MegatronBertForQuestionAnswering,
          '''text-classification''': MegatronBertForSequenceClassification,
          '''text-generation''': MegatronBertForCausalLM,
          '''token-classification''': MegatronBertForTokenClassification,
          '''zero-shot''': MegatronBertForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  _lowerCamelCase    =  True
  # test_resize_embeddings = False
  _lowerCamelCase    =  False


  def  lowerCamelCase       (      self		:			Any       ,				lowerCamelCase		:			Optional[Any]       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			Optional[int]=False							)						->	Optional[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    super()._prepare_for_class(lowerCamelCase       ,				lowerCamelCase       ,				return_labels=lowerCamelCase							)

         if return_labels:
                if model_class in get_values(lowerCamelCase							):
                       _UpperCAmelCase							=    torch.zeros(
                           (self.model_tester.batch_size, self.model_tester.seq_length)       ,				dtype=torch.long       ,				device=lowerCamelCase							)
                       _UpperCAmelCase							=    torch.zeros(
                           self.model_tester.batch_size       ,				dtype=torch.long       ,				device=lowerCamelCase							)
         return inputs_dict


  def  lowerCamelCase       (      self		:			Dict							)						->	Union[str, Any]:


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertModelTester(self							)
         _UpperCAmelCase							=    ConfigTester(self       ,				config_class=lowerCamelCase       ,				hidden_size=37							)


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


         """simple docstring"""
         self.config_tester.run_common_tests()


  def  lowerCamelCase       (      self		:			Union[str, Any]							)						->	List[str]:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_model(*lowerCamelCase							)


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


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			Dict							)						->	Optional[int]:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			List[Any]							)						->	Optional[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCamelCase							)


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


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			Tuple							)						->	Dict:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			List[str]							)						->	Optional[int]:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCamelCase							)





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


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCamelCase							)




def       _SCREAMING_SNAKE_CASE (							__snake_case )				->						Dict:
       return torch.tensor(
           __snake_case   ,      dtype=torch.long   ,      device=__snake_case   ,      )


__a:     int   					=					1E-4





@require_torch
@require_sentencepiece
@require_tokenizers
class 		SCREAMING_SNAKE_CASE__	(  unittest.TestCase     ):

  '''simple docstring'''







  @slow
  @unittest.skip("""Model is not available."""							)
  def  lowerCamelCase       (      self		:			int							)						->	Union[str, Any]:


         """simple docstring"""
         _UpperCAmelCase							=    """nvidia/megatron-bert-uncased-345m"""
         if "MYDIR" in os.environ:
                _UpperCAmelCase							=    os.path.join(os.environ["""MYDIR"""]       ,				lowerCamelCase							)
         _UpperCAmelCase							=    MegatronBertModel.from_pretrained(lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.half()
         _UpperCAmelCase							=    _long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]							)
         with torch.no_grad():
                _UpperCAmelCase							=    model(lowerCamelCase							)[0]
         _UpperCAmelCase							=    torch.Size((1, 9, 1024)							)
         self.assertEqual(output.shape       ,				lowerCamelCase							)

         _UpperCAmelCase							=    [-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
         for ii in range(3							):
                for jj in range(3							):
                       _UpperCAmelCase							=    output[0, ii, jj]
                       _UpperCAmelCase							=    expected[3 * ii + jj]
                       _UpperCAmelCase							=    """ii={} jj={} a={} b={}""".format(lowerCamelCase       ,				lowerCamelCase       ,				lowerCamelCase       ,				lowerCamelCase							)
                       self.assertTrue(math.isclose(lowerCamelCase       ,				lowerCamelCase       ,				rel_tol=lowerCamelCase       ,				abs_tol=lowerCamelCase							)       ,				msg=lowerCamelCase							)



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 						a__			(				A__				):
			if isinstance(A__,      collections.abc.Iterable				):
						return x
			return (x, x)







@require_flax
class 	__lowercase    :





							"""simple docstring"""







							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										pass






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										pass






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										pass






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					np.abs((a - b)       ).max()
										self.assertLessEqual(lowerCAmelCase__			,						lowerCAmelCase__			,						F'''Difference between torch and flax is {diff} (>= {tol}).'''       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__=None			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase__			,						lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     List[str]					=					FlaxVisionTextDualEncoderModel(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Tuple					=					model(input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )

										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			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__=None			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.get_vision_text_model(lowerCAmelCase__			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     str					=					{'vision_model': vision_model, 'text_model': text_model}
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					model(input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )

										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			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__=None			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     int					=					self.get_vision_text_model(lowerCAmelCase__			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					{'vision_model': vision_model, 'text_model': text_model}
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					model(input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					output[0]

										with tempfile.TemporaryDirectory() as tmpdirname:
													model.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					FlaxVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase__       )

													SCREAMING_SNAKE_CASE_		:     Tuple					=					model(input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					after_output[0]
													SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					np.amax(np.abs(out_a - out_a       )       )
													self.assertLessEqual(lowerCAmelCase__			,						1E-3       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__=None			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.get_vision_text_model(lowerCAmelCase__			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     List[str]					=					{'vision_model': vision_model, 'text_model': text_model}
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     List[str]					=					model(
										    input_ids=lowerCAmelCase__			,						pixel_values=lowerCAmelCase__			,						attention_mask=lowerCAmelCase__			,						output_attentions=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     List[str]					=					output.vision_model_output.attentions
										self.assertEqual(len(lowerCAmelCase__       )			,						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_		:     str					=					to_atuple(vision_model.config.image_size       )
										SCREAMING_SNAKE_CASE_		:     Any					=					to_atuple(vision_model.config.patch_size       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
										SCREAMING_SNAKE_CASE_		:     str					=					num_patches + 1
										self.assertEqual(vision_attentions[0].shape[-3:]			,						(vision_config.num_attention_heads, seq_len, seq_len)       )

										SCREAMING_SNAKE_CASE_		:     Tuple					=					output.text_model_output.attentions
										self.assertEqual(len(lowerCAmelCase__       )			,						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			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										pt_model.to(lowerCAmelCase__       )
										pt_model.eval()

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

										with torch.no_grad():
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					pt_model(**lowerCAmelCase__       ).to_tuple()

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

										# PT -> Flax
										with tempfile.TemporaryDirectory() as tmpdirname:
													pt_model.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					FlaxVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase__			,						from_pt=lowerCAmelCase__       )

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

										# Flax -> PT
										with tempfile.TemporaryDirectory() as tmpdirname:
													fx_model.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					VisionTextDualEncoderModel.from_pretrained(lowerCAmelCase__			,						from_flax=lowerCAmelCase__       )

										pt_model_loaded.to(lowerCAmelCase__       )
										pt_model_loaded.eval()

										with torch.no_grad():
													SCREAMING_SNAKE_CASE_		:     Optional[int]					=					pt_model_loaded(**lowerCAmelCase__       ).to_tuple()

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






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[Any]					=					VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase__			,						lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					VisionTextDualEncoderModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					FlaxVisionTextDualEncoderModel(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     str					=					convert_pytorch_state_dict_to_flax(pt_model.state_dict()			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					fx_state

										self.check_pt_flax_equivalence(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase__			,						lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					VisionTextDualEncoderModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     str					=					FlaxVisionTextDualEncoderModel(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     int					=					load_flax_weights_in_pytorch_model(lowerCAmelCase__			,						fx_model.params       )

										self.check_pt_flax_equivalence(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     str					=					self.prepare_config_and_inputs()
										self.check_model_from_pretrained_configs(**lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					self.prepare_config_and_inputs()
										self.check_vision_text_dual_encoder_from_pretrained(**lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					self.prepare_config_and_inputs()
										self.check_save_load(**lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					self.prepare_config_and_inputs()
										self.check_vision_text_output_attention(**lowerCAmelCase__       )






							@is_pt_flax_cross_test
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					self.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					config_inputs_dict.pop('vision_config'       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					config_inputs_dict.pop('text_config'       )

										SCREAMING_SNAKE_CASE_		:     Any					=					config_inputs_dict

										self.check_equivalence_pt_to_flax(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )
										self.check_equivalence_flax_to_pt(lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__       )






							@slow
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Dict					=					self.get_pretrained_model_and_inputs()

										SCREAMING_SNAKE_CASE_		:     int					=					model_a(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					outputs[0]

										with tempfile.TemporaryDirectory() as tmp_dirname:
													model_a.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     int					=					FlaxVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase__       )

													SCREAMING_SNAKE_CASE_		:     str					=					model_a(**lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     Any					=					after_outputs[0]
													SCREAMING_SNAKE_CASE_		:     str					=					np.amax(np.abs(out_a - out_a       )       )
													self.assertLessEqual(lowerCAmelCase__			,						1E-5       )







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





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
										    'hf-internal-testing/tiny-random-vit'			,						'hf-internal-testing/tiny-bert'			,						vision_from_pt=lowerCAmelCase__			,						text_from_pt=lowerCAmelCase__			,						)
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					1_3
										SCREAMING_SNAKE_CASE_		:     str					=					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_		:     Tuple					=					ids_tensor([batch_size, 4]			,						model.config.text_config.vocab_size       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					random_attention_mask([batch_size, 4]       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					{'pixel_values': pixel_values, 'input_ids': input_ids, 'attention_mask': attention_mask}

										return model, inputs






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     int					=					FlaxViTModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					FlaxBertModel(lowerCAmelCase__       )
										return vision_model, text_model






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Any					=					vision_config_and_inputs

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     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_torch
class 	__lowercase    (__SCREAMING_SNAKE_CASE	,       unittest.TestCase  ):





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     str					=					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
										    'hf-internal-testing/tiny-random-clip'			,						'hf-internal-testing/tiny-bert'			,						vision_from_pt=lowerCAmelCase__			,						text_from_pt=lowerCAmelCase__			,						)
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					1_3
										SCREAMING_SNAKE_CASE_		:     Dict					=					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_		:     int					=					ids_tensor([batch_size, 4]			,						model.config.text_config.vocab_size       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					random_attention_mask([batch_size, 4]       )
										SCREAMING_SNAKE_CASE_		:     int					=					{'pixel_values': pixel_values, 'input_ids': input_ids, 'attention_mask': attention_mask}

										return model, inputs






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					FlaxCLIPVisionModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Any					=					FlaxBertModel(lowerCAmelCase__       )
										return vision_model, text_model






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					FlaxCLIPVisionModelTester(self       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					FlaxBertModelTester(self       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					clip_model_tester.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE_		:     int					=					bert_model_tester.prepare_config_and_inputs()

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     int					=					vision_config_and_inputs

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     str					=					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  ):





							"""simple docstring"""







							@slow
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[str]					=					FlaxVisionTextDualEncoderModel.from_pretrained('clip-italian/clip-italian'			,						logit_scale_init_value=1.0       )
										SCREAMING_SNAKE_CASE_		:     str					=					VisionTextDualEncoderProcessor.from_pretrained('clip-italian/clip-italian'       )

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

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					model(**lowerCAmelCase__       )

										# 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_		:     Tuple					=					np.array([[1.2_284_727, 0.3_104_122]]       )

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





