


"""simple docstring"""



from __future__ import annotations

import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_tf, slow
from transformers.utils import is_tf_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask


if is_tf_available():
  import tensorflow as tf

  from transformers import TFBlipTextModel
  from transformers.models.blip.modeling_tf_blip import TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST







class _lowerCAmelCase       :







   def __init__(			self		,			UpperCamelCase__		,			UpperCamelCase__=12		,			UpperCamelCase__=7		,			UpperCamelCase__=True		,			UpperCamelCase__=True		,			UpperCamelCase__=True		,			UpperCamelCase__=99		,			UpperCamelCase__=32		,			UpperCamelCase__=32		,			UpperCamelCase__=2		,			UpperCamelCase__=4		,			UpperCamelCase__=37		,			UpperCamelCase__=0.1		,			UpperCamelCase__=0.1		,			UpperCamelCase__=512		,			UpperCamelCase__=0.02		,			UpperCamelCase__=0		,			UpperCamelCase__=None		,			)			->	Union[str, Any]:

       '''simple docstring'''




       snake_case					:   str         =					parent
       snake_case					:   Union[str, Any]         =					batch_size
       snake_case					:   int         =					seq_length
       snake_case					:   Dict         =					is_training
       snake_case					:   Optional[Any]         =					use_input_mask
       snake_case					:   Optional[Any]         =					use_labels
       snake_case					:   int         =					vocab_size
       snake_case					:   Union[str, Any]         =					hidden_size
       snake_case					:   Dict         =					projection_dim
       snake_case					:   str         =					num_hidden_layers
       snake_case					:   List[str]         =					num_attention_heads
       snake_case					:   List[str]         =					intermediate_size
       snake_case					:   int         =					dropout
       snake_case					:   Dict         =					attention_dropout
       snake_case					:   Tuple         =					max_position_embeddings
       snake_case					:   List[str]         =					initializer_range
       snake_case					:   Any         =					scope
       snake_case					:   List[Any]         =					bos_token_id







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   List[Any]         =					ids_tensor([self.batch_size, self.seq_length]		,			self.vocab_size		)

       snake_case					:   Optional[Any]         =					None
       if self.use_input_mask:
           snake_case					:   Optional[Any]         =					random_attention_mask([self.batch_size, self.seq_length]		)

       if input_mask is not None:
           snake_case					:   Dict         =					input_mask.numpy()
           snake_case      ,snake_case					:   List[str]         =					input_mask.shape
           snake_case					:   Union[str, Any]         =					np.random.randint(1		,			seq_length - 1		,			size=(batch_size,)		)
           for batch_idx, start_index in enumerate(UpperCamelCase__		):
               snake_case					:   Optional[Any]         =					1
               snake_case					:   Optional[Any]         =					0

       snake_case					:   Optional[Any]         =					self.get_config()

       return config, input_ids, tf.convert_to_tensor(UpperCamelCase__		)







   def 				lowerCamelCase	(			self		)			->	Dict:

       '''simple docstring'''




       return BlipTextConfig(
           vocab_size=self.vocab_size		,			hidden_size=self.hidden_size		,			projection_dim=self.projection_dim		,			num_hidden_layers=self.num_hidden_layers		,			num_attention_heads=self.num_attention_heads		,			intermediate_size=self.intermediate_size		,			dropout=self.dropout		,			attention_dropout=self.attention_dropout		,			max_position_embeddings=self.max_position_embeddings		,			initializer_range=self.initializer_range		,			bos_token_id=self.bos_token_id		,			)







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

       '''simple docstring'''




       snake_case					:   Union[str, Any]         =					TFBlipTextModel(config=UpperCamelCase__		)
       snake_case					:   Any         =					model(UpperCamelCase__		,			attention_mask=UpperCamelCase__		,			training=UpperCamelCase__		)
       snake_case					:   Optional[Any]         =					model(UpperCamelCase__		,			training=UpperCamelCase__		)
       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		)			->	int:

       '''simple docstring'''




       snake_case					:   Union[str, Any]         =					self.prepare_config_and_inputs()
       snake_case      ,snake_case      ,snake_case					:   Union[str, Any]         =					config_and_inputs
       snake_case					:   List[Any]         =					{"input_ids": input_ids, "attention_mask": input_mask}
       return config, inputs_dict







@require_tf
class _lowerCAmelCase       (   snake_case_					,							unittest.TestCase			):
   __UpperCAmelCase				: List[str]													=						(TFBlipTextModel,) if is_tf_available() else ()
   __UpperCAmelCase				: int													=						False
   __UpperCAmelCase				: List[Any]													=						False
   __UpperCAmelCase				: List[str]													=						False







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

       '''simple docstring'''




       snake_case					:   Dict         =					BlipTextModelTester(self		)
       snake_case					:   Union[str, Any]         =					ConfigTester(self		,			config_class=UpperCamelCase__		,			hidden_size=37		)







   def 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




       self.config_tester.run_common_tests()







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

       '''simple docstring'''




       snake_case					:   str         =					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*UpperCamelCase__		)







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

       '''simple docstring'''




       pass







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

       '''simple docstring'''




       pass







   @unittest.skip(reason="Blip does not use inputs_embeds"		)
   def 				lowerCamelCase	(			self		)			->	List[str]:

       '''simple docstring'''




       pass







   @unittest.skip(reason="BlipTextModel has no base class and is not available in MODEL_MAPPING"		)
   def 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




       pass







   @unittest.skip(reason="BlipTextModel has no base class and is not available in MODEL_MAPPING"		)
   def 				lowerCamelCase	(			self		)			->	int:

       '''simple docstring'''




       pass







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

       '''simple docstring'''




       for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           snake_case					:   Any         =					TFBlipTextModel.from_pretrained(UpperCamelCase__		)
           self.assertIsNotNone(UpperCamelCase__		)







   def 				lowerCamelCase	(			self		,			UpperCamelCase__=True		)			->	List[Any]:

       '''simple docstring'''




       super().test_pt_tf_model_equivalence(allow_missing_keys=UpperCamelCase__		)









"""simple docstring"""



import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   Optional[int]         =					[]
       snake_case					:   Tuple         =					[]

       for i in range(self.num_layers		):
           snake_case					:   Any         =					self.in_channels if i == 0 else self.out_channels

           snake_case					:   Optional[int]         =					FlaxResnetBlockaD(
               in_channels=UpperCamelCase__		,			out_channels=self.out_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)

           snake_case					:   Dict         =					FlaxTransformeraDModel(
               in_channels=self.out_channels		,			n_heads=self.num_attention_heads		,			d_head=self.out_channels // self.num_attention_heads		,			depth=1		,			use_linear_projection=self.use_linear_projection		,			only_cross_attention=self.only_cross_attention		,			use_memory_efficient_attention=self.use_memory_efficient_attention		,			dtype=self.dtype		,			)
           attentions.append(UpperCamelCase__		)

       snake_case					:   List[Any]         =					resnets
       snake_case					:   List[Any]         =					attentions

       if self.add_downsample:
           snake_case					:   Optional[Any]         =					FlaxDownsampleaD(self.out_channels		,			dtype=self.dtype		)







   def __call__(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__=True		)			->	List[str]:

       '''simple docstring'''




       snake_case					:   Tuple         =					()

       for resnet, attn in zip(self.resnets		,			self.attentions		):
           snake_case					:   str         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           snake_case					:   int         =					attn(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           output_states += (hidden_states,)

       if self.add_downsample:
           snake_case					:   Any         =					self.downsamplers_a(UpperCamelCase__		)
           output_states += (hidden_states,)

       return hidden_states, output_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   Dict         =					[]

       for i in range(self.num_layers		):
           snake_case					:   List[str]         =					self.in_channels if i == 0 else self.out_channels

           snake_case					:   str         =					FlaxResnetBlockaD(
               in_channels=UpperCamelCase__		,			out_channels=self.out_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)
       snake_case					:   str         =					resnets

       if self.add_downsample:
           snake_case					:   str         =					FlaxDownsampleaD(self.out_channels		,			dtype=self.dtype		)







   def __call__(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__=True		)			->	List[Any]:

       '''simple docstring'''




       snake_case					:   Tuple         =					()

       for resnet in self.resnets:
           snake_case					:   Optional[Any]         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           output_states += (hidden_states,)

       if self.add_downsample:
           snake_case					:   Tuple         =					self.downsamplers_a(UpperCamelCase__		)
           output_states += (hidden_states,)

       return hidden_states, output_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   Dict         =					[]
       snake_case					:   Union[str, Any]         =					[]

       for i in range(self.num_layers		):
           snake_case					:   Optional[Any]         =					self.in_channels if (i == self.num_layers - 1) else self.out_channels
           snake_case					:   Union[str, Any]         =					self.prev_output_channel if i == 0 else self.out_channels

           snake_case					:   Optional[Any]         =					FlaxResnetBlockaD(
               in_channels=resnet_in_channels + res_skip_channels		,			out_channels=self.out_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)

           snake_case					:   int         =					FlaxTransformeraDModel(
               in_channels=self.out_channels		,			n_heads=self.num_attention_heads		,			d_head=self.out_channels // self.num_attention_heads		,			depth=1		,			use_linear_projection=self.use_linear_projection		,			only_cross_attention=self.only_cross_attention		,			use_memory_efficient_attention=self.use_memory_efficient_attention		,			dtype=self.dtype		,			)
           attentions.append(UpperCamelCase__		)

       snake_case					:   str         =					resnets
       snake_case					:   List[Any]         =					attentions

       if self.add_upsample:
           snake_case					:   str         =					FlaxUpsampleaD(self.out_channels		,			dtype=self.dtype		)







   def __call__(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__=True		)			->	Optional[int]:

       '''simple docstring'''




       for resnet, attn in zip(self.resnets		,			self.attentions		):
           # pop res hidden states
           snake_case					:   Dict         =					res_hidden_states_tuple[-1]
           snake_case					:   Optional[int]         =					res_hidden_states_tuple[:-1]
           snake_case					:   List[Any]         =					jnp.concatenate((hidden_states, res_hidden_states)		,			axis=-1		)

           snake_case					:   List[Any]         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           snake_case					:   Any         =					attn(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)

       if self.add_upsample:
           snake_case					:   int         =					self.upsamplers_a(UpperCamelCase__		)

       return hidden_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   str         =					[]

       for i in range(self.num_layers		):
           snake_case					:   Union[str, Any]         =					self.in_channels if (i == self.num_layers - 1) else self.out_channels
           snake_case					:   Optional[Any]         =					self.prev_output_channel if i == 0 else self.out_channels

           snake_case					:   Dict         =					FlaxResnetBlockaD(
               in_channels=resnet_in_channels + res_skip_channels		,			out_channels=self.out_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)

       snake_case					:   Optional[int]         =					resnets

       if self.add_upsample:
           snake_case					:   Dict         =					FlaxUpsampleaD(self.out_channels		,			dtype=self.dtype		)







   def __call__(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__=True		)			->	int:

       '''simple docstring'''




       for resnet in self.resnets:
           # pop res hidden states
           snake_case					:   int         =					res_hidden_states_tuple[-1]
           snake_case					:   List[Any]         =					res_hidden_states_tuple[:-1]
           snake_case					:   Optional[int]         =					jnp.concatenate((hidden_states, res_hidden_states)		,			axis=-1		)

           snake_case					:   Dict         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)

       if self.add_upsample:
           snake_case					:   Optional[int]         =					self.upsamplers_a(UpperCamelCase__		)

       return hidden_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







   def 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




       snake_case					:   Dict         =					[
           FlaxResnetBlockaD(
               in_channels=self.in_channels		,			out_channels=self.in_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
       ]

       snake_case					:   Any         =					[]

       for _ in range(self.num_layers		):
           snake_case					:   Union[str, Any]         =					FlaxTransformeraDModel(
               in_channels=self.in_channels		,			n_heads=self.num_attention_heads		,			d_head=self.in_channels // self.num_attention_heads		,			depth=1		,			use_linear_projection=self.use_linear_projection		,			use_memory_efficient_attention=self.use_memory_efficient_attention		,			dtype=self.dtype		,			)
           attentions.append(UpperCamelCase__		)

           snake_case					:   List[Any]         =					FlaxResnetBlockaD(
               in_channels=self.in_channels		,			out_channels=self.in_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)

       snake_case					:   Union[str, Any]         =					resnets
       snake_case					:   str         =					attentions







   def __call__(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__=True		)			->	Any:

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					self.resnets[0](UpperCamelCase__		,			UpperCamelCase__		)
       for attn, resnet in zip(self.attentions		,			self.resnets[1:]		):
           snake_case					:   Union[str, Any]         =					attn(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           snake_case					:   Optional[Any]         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)

       return hidden_states





