




import collections
import inspect
import unittest

from transformers import NatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch
							from torch import nn

							from transformers import NatBackbone, NatForImageClassification, NatModel
							from transformers.models.nat.modeling_nat import NAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor


class 				__SCREAMING_SNAKE_CASE							:

					def __init__(				self ,       __lowercase ,       __lowercase=13 ,       __lowercase=64 ,       __lowercase=4 ,       __lowercase=3 ,       __lowercase=16 ,       __lowercase=[1, 2, 1] ,       __lowercase=[2, 4, 8] ,       __lowercase=3 ,       __lowercase=2.0 ,       __lowercase=True ,       __lowercase=0.0 ,       __lowercase=0.0 ,       __lowercase=0.1 ,       __lowercase="gelu" ,       __lowercase=True ,       __lowercase=0.02 ,       __lowercase=1E-5 ,       __lowercase=True ,       __lowercase=None ,       __lowercase=True ,       __lowercase=10 ,       __lowercase=["stage1", "stage2"] ,       __lowercase=[1, 2] ,       ):







						'''simple docstring'''




						lowerCamelCase_							:		int							   =							parent
						lowerCamelCase_							:		Tuple							   =							batch_size
						lowerCamelCase_							:		List[Any]							   =							image_size
						lowerCamelCase_							:		Union[str, Any]							   =							patch_size
						lowerCamelCase_							:		Any							   =							num_channels
						lowerCamelCase_							:		Tuple							   =							embed_dim
						lowerCamelCase_							:		Dict							   =							depths
						lowerCamelCase_							:		Optional[int]							   =							num_heads
						lowerCamelCase_							:		Tuple							   =							kernel_size
						lowerCamelCase_							:		Optional[int]							   =							mlp_ratio
						lowerCamelCase_							:		Any							   =							qkv_bias
						lowerCamelCase_							:		List[str]							   =							hidden_dropout_prob
						lowerCamelCase_							:		Dict							   =							attention_probs_dropout_prob
						lowerCamelCase_							:		Union[str, Any]							   =							drop_path_rate
						lowerCamelCase_							:		Optional[Any]							   =							hidden_act
						lowerCamelCase_							:		Optional[int]							   =							patch_norm
						lowerCamelCase_							:		Dict							   =							layer_norm_eps
						lowerCamelCase_							:		str							   =							initializer_range
						lowerCamelCase_							:		Optional[int]							   =							is_training
						lowerCamelCase_							:		int							   =							scope
						lowerCamelCase_							:		Optional[int]							   =							use_labels
						lowerCamelCase_							:		List[Any]							   =							num_labels
						lowerCamelCase_							:		Tuple							   =							out_features
						lowerCamelCase_							:		Optional[int]							   =							out_indices

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		str							   =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

						lowerCamelCase_							:		int							   =							None
						if self.use_labels:
							lowerCamelCase_							:		Optional[Any]							   =							ids_tensor([self.batch_size] ,       self.num_labels					)

						lowerCamelCase_							:		int							   =							self.get_config()

						return config, pixel_values, labels

					def 	lowercase_							(				self					):







						'''simple docstring'''




						return NatConfig(
						    num_labels=self.num_labels ,       image_size=self.image_size ,       patch_size=self.patch_size ,       num_channels=self.num_channels ,       embed_dim=self.embed_dim ,       depths=self.depths ,       num_heads=self.num_heads ,       kernel_size=self.kernel_size ,       mlp_ratio=self.mlp_ratio ,       qkv_bias=self.qkv_bias ,       hidden_dropout_prob=self.hidden_dropout_prob ,       attention_probs_dropout_prob=self.attention_probs_dropout_prob ,       drop_path_rate=self.drop_path_rate ,       hidden_act=self.hidden_act ,       patch_norm=self.patch_norm ,       layer_norm_eps=self.layer_norm_eps ,       initializer_range=self.initializer_range ,       out_features=self.out_features ,       out_indices=self.out_indices ,       )

					def 	lowercase_							(				self ,       __lowercase ,       __lowercase ,       __lowercase					):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[int]							   =							NatModel(config=__lowercase					)
						model.to(__lowercase					)
						model.eval()
						lowerCamelCase_							:		Tuple							   =							model(__lowercase					)

						lowerCamelCase_							:		Dict							   =							(config.image_size // config.patch_size) // (2 ** (len(config.depths					) - 1))
						lowerCamelCase_							:		Optional[int]							   =							int(config.embed_dim * 2 ** (len(config.depths					) - 1)					)

						self.parent.assertEqual(
						    result.last_hidden_state.shape ,       (self.batch_size, expected_height, expected_width, expected_dim)					)

					def 	lowercase_							(				self ,       __lowercase ,       __lowercase ,       __lowercase					):







						'''simple docstring'''




						lowerCamelCase_							:		Any							   =							NatForImageClassification(__lowercase					)
						model.to(__lowercase					)
						model.eval()
						lowerCamelCase_							:		int							   =							model(__lowercase ,       labels=__lowercase					)
						self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_labels)					)

						# test greyscale images
						lowerCamelCase_							:		List[str]							   =							1
						lowerCamelCase_							:		Dict							   =							NatForImageClassification(__lowercase					)
						model.to(__lowercase					)
						model.eval()

						lowerCamelCase_							:		Any							   =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
						lowerCamelCase_							:		Any							   =							model(__lowercase					)
						self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_labels)					)

					def 	lowercase_							(				self ,       __lowercase ,       __lowercase ,       __lowercase					):







						'''simple docstring'''




						lowerCamelCase_							:		Tuple							   =							NatBackbone(config=__lowercase					)
						model.to(__lowercase					)
						model.eval()
						lowerCamelCase_							:		Union[str, Any]							   =							model(__lowercase					)

						# verify hidden states
						self.parent.assertEqual(len(result.feature_maps					) ,       len(config.out_features					)					)
						self.parent.assertListEqual(list(result.feature_maps[0].shape					) ,       [self.batch_size, model.channels[0], 16, 16]					)

						# verify channels
						self.parent.assertEqual(len(model.channels					) ,       len(config.out_features					)					)

						# verify backbone works with out_features=None
						lowerCamelCase_							:		Optional[Any]							   =							None
						lowerCamelCase_							:		str							   =							NatBackbone(config=__lowercase					)
						model.to(__lowercase					)
						model.eval()
						lowerCamelCase_							:		Tuple							   =							model(__lowercase					)

						# verify feature maps
						self.parent.assertEqual(len(result.feature_maps					) ,       1					)
						self.parent.assertListEqual(list(result.feature_maps[0].shape					) ,       [self.batch_size, model.channels[-1], 4, 4]					)

						# verify channels
						self.parent.assertEqual(len(model.channels					) ,       1					)





					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		int							   =							self.prepare_config_and_inputs()
						lowerCamelCase_   ,			lowerCamelCase_   ,			lowerCamelCase_							:		List[Any]							   =							config_and_inputs
						lowerCamelCase_							:		Dict							   =							{'pixel_values': pixel_values}
						return config, inputs_dict


@require_natten
@require_torch
class 				__SCREAMING_SNAKE_CASE							(		lowerCamelCase   ,     lowerCamelCase   ,     unittest.TestCase  ):
					_lowerCamelCase		  = (
					    (
					        NatModel,
					        NatForImageClassification,
					        NatBackbone,
					    )
					    if is_torch_available()
					    else ()
					)
					_lowerCamelCase		  = (
					    {"""feature-extraction""": NatModel, """image-classification""": NatForImageClassification}
					    if is_torch_available()
					    else {}
					)
					_lowerCamelCase		  = False

					_lowerCamelCase		  = False
					_lowerCamelCase		  = False
					_lowerCamelCase		  = False
					_lowerCamelCase		  = False

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		List[Any]							   =							NatModelTester(self					)
						lowerCamelCase_							:		Optional[Any]							   =							ConfigTester(self ,       config_class=__lowercase ,       embed_dim=37					)

					def 	lowercase_							(				self					):







						'''simple docstring'''




						self.create_and_test_config_common_properties()
						self.config_tester.create_and_test_config_to_json_string()
						self.config_tester.create_and_test_config_to_json_file()
						self.config_tester.create_and_test_config_from_and_save_pretrained()
						self.config_tester.create_and_test_config_with_num_labels()
						self.config_tester.check_config_can_be_init_without_params()
						self.config_tester.check_config_arguments_init()

					def 	lowercase_							(				self					):







						'''simple docstring'''




						return

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		Any							   =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__lowercase					)

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		List[Any]							   =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*__lowercase					)

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		Union[str, Any]							   =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_backbone(*__lowercase					)

					@unittest.skip(reason='Nat does not use inputs_embeds'					)
					def 	lowercase_							(				self					):







						'''simple docstring'''




						pass

					@unittest.skip(reason='Nat does not use feedforward chunking'					)
					def 	lowercase_							(				self					):







						'''simple docstring'''




						pass

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_   ,			lowerCamelCase_							:		Tuple							   =							self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
							lowerCamelCase_							:		Any							   =							model_class(__lowercase					)
							self.assertIsInstance(model.get_input_embeddings() ,       (nn.Module)					)
							lowerCamelCase_							:		List[Any]							   =							model.get_output_embeddings()
							self.assertTrue(x is None or isinstance(__lowercase ,       nn.Linear					)					)

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_   ,			lowerCamelCase_							:		Dict							   =							self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
							lowerCamelCase_							:		Union[str, Any]							   =							model_class(__lowercase					)
							lowerCamelCase_							:		Dict							   =							inspect.signature(model.forward					)
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							lowerCamelCase_							:		int							   =							[*signature.parameters.keys()]

							lowerCamelCase_							:		List[Any]							   =							['pixel_values']
							self.assertListEqual(arg_names[:1] ,       __lowercase					)

					def 	lowercase_							(				self					):







						'''simple docstring'''




						self.skipTest('Nat\'s attention operation is handled entirely by NATTEN.'					)

					def 	lowercase_							(				self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase					):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[Any]							   =							model_class(__lowercase					)
						model.to(__lowercase					)
						model.eval()

						with torch.no_grad():
							lowerCamelCase_							:		Optional[int]							   =							model(**self._prepare_for_class(__lowercase ,       __lowercase					)					)

						lowerCamelCase_							:		Tuple							   =							outputs.hidden_states

						lowerCamelCase_							:		str							   =							getattr(
						    self.model_tester ,       'expected_num_hidden_layers' ,       len(self.model_tester.depths					) + 1					)
						self.assertEqual(len(__lowercase					) ,       __lowercase					)

						# Nat has a different seq_length
						lowerCamelCase_							:		Any							   =							(
						    config.patch_size
						    if isinstance(config.patch_size ,       collections.abc.Iterable					)
						    else (config.patch_size, config.patch_size)
						)

						lowerCamelCase_							:		List[Any]							   =							image_size[0] // patch_size[0]
						lowerCamelCase_							:		Any							   =							image_size[1] // patch_size[1]

						self.assertListEqual(
						    list(hidden_states[0].shape[-3:]					) ,       [height, width, self.model_tester.embed_dim] ,       )

						if model_class.__name__ != "NatBackbone":
							lowerCamelCase_							:		Dict							   =							outputs.reshaped_hidden_states
							self.assertEqual(len(__lowercase					) ,       __lowercase					)

							lowerCamelCase_   ,			lowerCamelCase_   ,			lowerCamelCase_   ,			lowerCamelCase_							:		Optional[int]							   =							reshaped_hidden_states[0].shape
							lowerCamelCase_							:		Dict							   =							(
							    reshaped_hidden_states[0].view(__lowercase ,       __lowercase ,       __lowercase ,       __lowercase					).permute(0 ,       2 ,       3 ,       1					)
							)
							self.assertListEqual(
							    list(reshaped_hidden_states.shape[-3:]					) ,       [height, width, self.model_tester.embed_dim] ,       )

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_   ,			lowerCamelCase_							:		str							   =							self.model_tester.prepare_config_and_inputs_for_common()

						lowerCamelCase_							:		List[Any]							   =							(
						    self.model_tester.image_size
						    if isinstance(self.model_tester.image_size ,       collections.abc.Iterable					)
						    else (self.model_tester.image_size, self.model_tester.image_size)
						)

						for model_class in self.all_model_classes:
							lowerCamelCase_							:		str							   =							True
							self.check_hidden_states_output(__lowercase ,       __lowercase ,       __lowercase ,       __lowercase					)

							# check that output_hidden_states also work using config
							del inputs_dict["output_hidden_states"]
							lowerCamelCase_							:		Union[str, Any]							   =							True

							self.check_hidden_states_output(__lowercase ,       __lowercase ,       __lowercase ,       __lowercase					)

					@slow
					def 	lowercase_							(				self					):







						'''simple docstring'''




						for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							lowerCamelCase_							:		Optional[Any]							   =							NatModel.from_pretrained(__lowercase					)
							self.assertIsNotNone(__lowercase					)





					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_   ,			lowerCamelCase_							:		str							   =							self.model_tester.prepare_config_and_inputs_for_common()

						lowerCamelCase_							:		str							   =							_config_zero_init(__lowercase					)
						for model_class in self.all_model_classes:
							lowerCamelCase_							:		int							   =							model_class(config=__lowercase					)
							for name, param in model.named_parameters():
								if "embeddings" not in name and param.requires_grad:
									self.assertIn(
									    ((param.data.mean() * 1E9).round() / 1E9).item() ,       [0.0, 1.0] ,       msg=f"Parameter {name} of model {model_class} seems not properly initialized" ,       )


@require_natten
@require_vision
@require_torch
class 				__SCREAMING_SNAKE_CASE							(		unittest.TestCase  ):

					@cached_property
					def 	lowercase_							(				self					):







						'''simple docstring'''




						return AutoImageProcessor.from_pretrained('shi-labs/nat-mini-in1k-224'					) if is_vision_available() else None





					@slow
					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		str							   =							NatForImageClassification.from_pretrained('shi-labs/nat-mini-in1k-224'					).to(__lowercase					)
						lowerCamelCase_							:		int							   =							self.default_image_processor

						lowerCamelCase_							:		Any							   =							Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'					)
						lowerCamelCase_							:		Any							   =							image_processor(images=__lowercase ,       return_tensors='pt'					).to(__lowercase					)

						# forward pass
						with torch.no_grad():
							lowerCamelCase_							:		Any							   =							model(**__lowercase					)

						# verify the logits
						lowerCamelCase_							:		Union[str, Any]							   =							torch.Size((1, 1_000)					)
						self.assertEqual(outputs.logits.shape ,       __lowercase					)
						lowerCamelCase_							:		List[Any]							   =							torch.tensor([0.38_05, -0.86_76, -0.39_12]					).to(__lowercase					)
						self.assertTrue(torch.allclose(outputs.logits[0, :3] ,       __lowercase ,       atol=1E-4					)					)


@require_torch
@require_natten
class 				__SCREAMING_SNAKE_CASE							(		unittest.TestCase   ,     lowerCamelCase  ):
					_lowerCamelCase		  = (NatBackbone,) if is_torch_available() else ()
					_lowerCamelCase		  = NatConfig

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		List[str]							   =							NatModelTester(self					)




from typing import Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from configuration_hybrid_clip import HybridCLIPConfig
from flax.core.frozen_dict import FrozenDict

from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
from transformers.modeling_flax_utils import FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPOutput
from transformers.utils import logging


lowercase_           =	logging.get_logger(__name__)


class        A_					(       nn.Module    ):




      '''simple docstring'''


      __snake_case					       =   42
      __snake_case					       =   jnp.floataa




      def      _snake_case	(				self:       List[Any]				):
       __lowerCamelCase		:							Union[str, Any]							=						self.config.text_config
       __lowerCamelCase		:							List[Any]							=						self.config.vision_config

       __lowerCamelCase		:							Any							=						self.config.projection_dim
       __lowerCamelCase		:							Dict							=						text_config.hidden_size
       __lowerCamelCase		:							str							=						vision_config.hidden_size

       __lowerCamelCase		:							Optional[Any]							=						FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
       __lowerCamelCase		:							int							=						FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__						,      a				).module_class

       __lowerCamelCase		:							Union[str, Any]							=						text_module(a						,      dtype=self.dtype				)
       __lowerCamelCase		:							int							=						vision_module(a						,      dtype=self.dtype				)

       __lowerCamelCase		:							Union[str, Any]							=						nn.Dense(
           self.projection_dim						,      dtype=self.dtype						,      kernel_init=jax.nn.initializers.normal(0.0_2				)						,      use_bias=a						,      )
       __lowerCamelCase		:							int							=						nn.Dense(
           self.projection_dim						,      dtype=self.dtype						,      kernel_init=jax.nn.initializers.normal(0.0_2				)						,      use_bias=a						,      )
       __lowerCamelCase		:							Dict							=						self.param('logit_scale'						,      jax.nn.initializers.ones						,      []				)




      def __call__(				self:       int						,      a:       Union[str, Any]=None						,      a:       str=None						,      a:       Tuple=None						,      a:       str=None						,      a:       List[Any]=None						,      a:       bool = True						,      a:       List[str]=None						,      a:       Any=None						,      a:       Optional[Any]=None						,      ):
       __lowerCamelCase		:							str							=						return_dict if return_dict is not None else self.config.return_dict

       __lowerCamelCase		:							Any							=						self.vision_model(
           pixel_values=a						,      deterministic=a						,      output_attentions=a						,      output_hidden_states=a						,      return_dict=a						,      )

       __lowerCamelCase		:							Union[str, Any]							=						self.text_model(
           input_ids=a						,      attention_mask=a						,      token_type_ids=a						,      position_ids=a						,      deterministic=a						,      output_attentions=a						,      output_hidden_states=a						,      return_dict=a						,      )

       __lowerCamelCase		:							Tuple							=						vision_outputs[1]
       __lowerCamelCase		:							List[Any]							=						self.visual_projection(a				)

       __lowerCamelCase		:							Any							=						text_outputs[1]
       __lowerCamelCase		:							Optional[Any]							=						self.text_projection(a				)

       # normalized features
       __lowerCamelCase		:							str							=						image_embeds / jnp.linalg.norm(a						,      axis=-1						,      keepdims=a				)
       __lowerCamelCase		:							Optional[int]							=						text_embeds / jnp.linalg.norm(a						,      axis=-1						,      keepdims=a				)

       # cosine similarity as logits
       __lowerCamelCase		:							Union[str, Any]							=						jnp.exp(self.logit_scale				)
       __lowerCamelCase		:							Union[str, Any]							=						jnp.matmul(a						,      image_embeds.T				) * logit_scale
       __lowerCamelCase		:							List[str]							=						logits_per_text.T

       if not return_dict:
        return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)

       return FlaxCLIPOutput(
           logits_per_image=a						,      logits_per_text=a						,      text_embeds=a						,      image_embeds=a						,      text_model_output=a						,      vision_model_output=a						,      )



class        A_					(       __UpperCamelCase    ):




      '''simple docstring'''


      __snake_case					       =   HybridCLIPConfig
      __snake_case					       =   FlaxHybridCLIPModule




      def __init__(				self:       List[Any]						,      a:       HybridCLIPConfig						,      a:       Optional[Tuple] = None						,      a:       int = 0						,      a:       jnp.dtype = jnp.floataa						,      **a:       int						,      ):
       if input_shape is None:
        __lowerCamelCase		:							int							=						((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

       __lowerCamelCase		:							int							=						self.module_class(config=a						,      dtype=a						,      **a				)
       super().__init__(a						,      a						,      input_shape=a						,      seed=a						,      dtype=a				)




      def      _snake_case	(				self:       Dict						,      a:       jax.random.PRNGKey						,      a:       Tuple						,      a:       FrozenDict = None				):
       # init input tensor
       __lowerCamelCase		:							int							=						jnp.zeros(input_shape[0]						,      dtype='i4'				)
       __lowerCamelCase		:							List[Any]							=						jnp.broadcast_to(jnp.arange(jnp.atleast_ad(a				).shape[-1]				)						,      input_shape[0]				)
       __lowerCamelCase		:							Any							=						jnp.ones_like(a				)
       __lowerCamelCase		:							int							=						jnp.ones_like(a				)

       __lowerCamelCase		:							Optional[Any]							=						jax.random.normal(a						,      input_shape[1]				)

       __lowerCamelCase					,     __lowerCamelCase		:							List[Any]							=						jax.random.split(a				)
       __lowerCamelCase		:							Tuple							=						{'params': params_rng, 'dropout': dropout_rng}

       return self.module.init(a						,      a						,      a						,      a						,      a						,      a				)["params"]




      def __call__(				self:       Tuple						,      a:       Optional[Any]						,      a:       Dict						,      a:       Optional[int]=None						,      a:       Union[str, Any]=None						,      a:       int=None						,      a:       dict = None						,      a:       jax.random.PRNGKey = None						,      a:       bool = False						,      a:       Optional[bool] = None						,      a:       Optional[bool] = None						,      a:       Optional[bool] = None						,      ):
       __lowerCamelCase		:							str							=						output_attentions if output_attentions is not None else self.config.output_attentions
       __lowerCamelCase		:							Dict							=						(
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )
       __lowerCamelCase		:							int							=						return_dict if return_dict is not None else self.config.return_dict

       if position_ids is None:
        __lowerCamelCase		:							Optional[int]							=						jnp.broadcast_to(jnp.arange(jnp.atleast_ad(a				).shape[-1]				)						,      input_ids.shape				)

       if token_type_ids is None:
        __lowerCamelCase		:							str							=						jnp.zeros_like(a				)

       if attention_mask is None:
        __lowerCamelCase		:							str							=						jnp.ones_like(a				)

       # Handle any PRNG if needed
       __lowerCamelCase		:							Union[str, Any]							=						{}
       if dropout_rng is not None:
        __lowerCamelCase		:							Any							=						dropout_rng

       return self.module.apply(
           {'params': params or self.params}						,      jnp.array(a						,      dtype='i4'				)						,      jnp.array(a						,      dtype=jnp.floataa				)						,      jnp.array(a						,      dtype='i4'				)						,      jnp.array(a						,      dtype='i4'				)						,      jnp.array(a						,      dtype='i4'				)						,      not train						,      a						,      a						,      a						,      rngs=a						,      )




      def      _snake_case	(				self:       Tuple						,      a:       Tuple						,      a:       List[Any]=None						,      a:       Dict=None						,      a:       int=None						,      a:       dict = None						,      a:       jax.random.PRNGKey = None						,      a:       Tuple=False						,      ):
       if position_ids is None:
        __lowerCamelCase		:							str							=						jnp.broadcast_to(jnp.arange(jnp.atleast_ad(a				).shape[-1]				)						,      input_ids.shape				)

       if token_type_ids is None:
        __lowerCamelCase		:							Union[str, Any]							=						jnp.zeros_like(a				)

       if attention_mask is None:
        __lowerCamelCase		:							int							=						jnp.ones_like(a				)

       # Handle any PRNG if needed
       __lowerCamelCase		:							Optional[int]							=						{}
       if dropout_rng is not None:
        __lowerCamelCase		:							List[str]							=						dropout_rng

       def _get_features(a:       str						,      a:       Any						,      a:       Optional[int]						,      a:       Optional[int]						,      a:       Union[str, Any]						,      a:       Optional[int]				):
        __lowerCamelCase		:							Tuple							=						module.text_model(
            input_ids=a						,      attention_mask=a						,      position_ids=a						,      token_type_ids=a						,      deterministic=a						,      )
        __lowerCamelCase		:							int							=						text_outputs[1]
        __lowerCamelCase		:							Optional[Any]							=						module.text_projection(a				)
        return text_features

       return self.module.apply(
           {'params': params or self.params}						,      jnp.array(a						,      dtype='i4'				)						,      jnp.array(a						,      dtype='i4'				)						,      jnp.array(a						,      dtype='i4'				)						,      jnp.array(a						,      dtype='i4'				)						,      not train						,      method=_get_features						,      rngs=a						,      )




      def      _snake_case	(				self:       int						,      a:       Optional[int]						,      a:       dict = None						,      a:       jax.random.PRNGKey = None						,      a:       str=False				):
       __lowerCamelCase		:							Dict							=						{}
       if dropout_rng is not None:
        __lowerCamelCase		:							List[str]							=						dropout_rng

       def _get_features(a:       str						,      a:       Tuple						,      a:       Any				):
        __lowerCamelCase		:							Dict							=						module.vision_model(pixel_values=a						,      deterministic=a				)
        __lowerCamelCase		:							int							=						vision_outputs[1]  # pooled_output
        __lowerCamelCase		:							int							=						module.visual_projection(a				)
        return image_features

       return self.module.apply(
           {'params': params or self.params}						,      jnp.array(a						,      dtype=jnp.floataa				)						,      not train						,      method=_get_features						,      rngs=a						,      )




      @classmethod
      def      _snake_case	(				cls:       str						,      a:       str = None						,      a:       str = None						,      *a:       Optional[Any]						,      **a:       List[str]						,      ):
       __lowerCamelCase		:							List[Any]							=						{
           argument[len('text_'				) :]: value for argument, value in kwargs.items() if argument.startswith('text_'				)
       }

       __lowerCamelCase		:							Optional[Any]							=						{
           argument[len('vision_'				) :]: value for argument, value in kwargs.items() if argument.startswith('vision_'				)
       }

       # remove text, vision kwargs from kwargs
       for key in kwargs_text.keys():
        del kwargs["text_" + key]
       for key in kwargs_vision.keys():
        del kwargs["vision_" + key]

       # Load and initialize the text and vision model
       __lowerCamelCase		:							Any							=						kwargs_text.pop('model'						,      a				)
       if text_model is None:
        assert (
            text_model_name_or_path is not None
        ), "If `model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
        from transformers import FlaxAutoModel

        if "config" not in kwargs_text:
         from transformers import AutoConfig

         __lowerCamelCase		:							Optional[int]							=						AutoConfig.from_pretrained(a				)
         __lowerCamelCase		:							Any							=						text_config

        __lowerCamelCase		:							Dict							=						FlaxAutoModel.from_pretrained(a						,      *a						,      **a				)

       __lowerCamelCase		:							Any							=						kwargs_vision.pop('model'						,      a				)
       if vision_model is None:
        assert (
            vision_model_name_or_path is not None
        ), "If `model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
        from transformers import FlaxAutoModel

        if "config" not in kwargs_vision:
         from transformers import AutoConfig

         __lowerCamelCase		:							Any							=						AutoConfig.from_pretrained(a				)
         __lowerCamelCase		:							int							=						vision_config

        __lowerCamelCase		:							str							=						FlaxAutoModel.from_pretrained(a						,      *a						,      **a				)

       # instantiate config with corresponding kwargs
       __lowerCamelCase		:							Optional[int]							=						kwargs.pop('dtype'						,      jnp.floataa				)
       __lowerCamelCase		:							List[Any]							=						HybridCLIPConfig.from_text_vision_configs(text_model.config						,      vision_model.config						,      **a				)

       # init model
       __lowerCamelCase		:							List[Any]							=						cls(a						,      *a						,      dtype=a						,      **a				)

       if vision_config.model_type == "clip":
        __lowerCamelCase		:							Any							=						vision_model.params['vision_model']
        __lowerCamelCase		:							Optional[Any]							=						vision_model.params['visual_projection']['kernel']
       else:
        __lowerCamelCase		:							Optional[int]							=						vision_model.params

       __lowerCamelCase		:							List[str]							=						text_model.params

       return model

