

"""simple docstring"""






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


lowerCamelCase__				:  Optional[int]													=  logging.get_logger(__name__)






class     _UpperCAmelCase			(		nn.Module):
			__a :       Optional[int]									=							4_2
			__a :       Optional[int]									=							jnp.floataa
			def 							__snake_case (		self     )    ->    Optional[Any]:






				'''simple docstring'''







				_UpperCAmelCase    :							Any									=					self.config.text_config
				_UpperCAmelCase    :							List[Any]									=					self.config.vision_config

				_UpperCAmelCase    :							str									=					self.config.projection_dim
				_UpperCAmelCase    :							int									=					text_config.hidden_size
				_UpperCAmelCase    :							Any									=					vision_config.hidden_size

				_UpperCAmelCase    :							str									=					FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
				_UpperCAmelCase    :							Any									=					FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__      ,		__UpperCamelCase     ).module_class

				_UpperCAmelCase    :							Any									=					text_module(__UpperCamelCase      ,		dtype=self.dtype     )
				_UpperCAmelCase    :							Optional[Any]									=					vision_module(__UpperCamelCase      ,		dtype=self.dtype     )

				_UpperCAmelCase    :							List[Any]									=					nn.Dense(
				    self.projection_dim      ,		dtype=self.dtype      ,		kernel_init=jax.nn.initializers.normal(0.02     )      ,		use_bias=__UpperCamelCase      ,		)
				_UpperCAmelCase    :							str									=					nn.Dense(
				    self.projection_dim      ,		dtype=self.dtype      ,		kernel_init=jax.nn.initializers.normal(0.02     )      ,		use_bias=__UpperCamelCase      ,		)
				_UpperCAmelCase    :							Optional[int]									=					self.param("""logit_scale"""      ,		jax.nn.initializers.ones      ,		[]     )
			def __call__(		self      ,		_A=None      ,		_A=None      ,		_A=None      ,		_A=None      ,		_A=None      ,		_A = True      ,		_A=None      ,		_A=None      ,		_A=None      ,		)    ->    Optional[Any]:






				'''simple docstring'''







				_UpperCAmelCase    :							Union[str, Any]									=					return_dict if return_dict is not None else self.config.return_dict

				_UpperCAmelCase    :							Optional[int]									=					self.vision_model(
				    pixel_values=__UpperCamelCase      ,		deterministic=__UpperCamelCase      ,		output_attentions=__UpperCamelCase      ,		output_hidden_states=__UpperCamelCase      ,		return_dict=__UpperCamelCase      ,		)

				_UpperCAmelCase    :							str									=					self.text_model(
				    input_ids=__UpperCamelCase      ,		attention_mask=__UpperCamelCase      ,		token_type_ids=__UpperCamelCase      ,		position_ids=__UpperCamelCase      ,		deterministic=__UpperCamelCase      ,		output_attentions=__UpperCamelCase      ,		output_hidden_states=__UpperCamelCase      ,		return_dict=__UpperCamelCase      ,		)

				_UpperCAmelCase    :							Optional[Any]									=					vision_outputs[1]
				_UpperCAmelCase    :							Any									=					self.visual_projection(__UpperCamelCase     )

				_UpperCAmelCase    :							Union[str, Any]									=					text_outputs[1]
				_UpperCAmelCase    :							int									=					self.text_projection(__UpperCamelCase     )

				# normalized features
				_UpperCAmelCase    :							Tuple									=					image_embeds / jnp.linalg.norm(__UpperCamelCase      ,		axis=-1      ,		keepdims=__UpperCamelCase     )
				_UpperCAmelCase    :							Optional[Any]									=					text_embeds / jnp.linalg.norm(__UpperCamelCase      ,		axis=-1      ,		keepdims=__UpperCamelCase     )

				# cosine similarity as logits
				_UpperCAmelCase    :							List[Any]									=					jnp.exp(self.logit_scale     )
				_UpperCAmelCase    :							Optional[int]									=					jnp.matmul(__UpperCamelCase      ,		image_embeds.T     ) * logit_scale
				_UpperCAmelCase    :							Any									=					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=__UpperCamelCase      ,		logits_per_text=__UpperCamelCase      ,		text_embeds=__UpperCamelCase      ,		image_embeds=__UpperCamelCase      ,		text_model_output=__UpperCamelCase      ,		vision_model_output=__UpperCamelCase      ,		)






class     _UpperCAmelCase			(		UpperCAmelCase_):
			__a :       Union[str, Any]									=							HybridCLIPConfig
			__a :       int									=							FlaxHybridCLIPModule
			def __init__(		self      ,		_A      ,		_A = None      ,		_A = 0      ,		_A = jnp.floataa      ,		**_A      ,		)    ->    Dict:






				'''simple docstring'''







				if input_shape is None:
					_UpperCAmelCase    :							List[Any]									=					((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

				_UpperCAmelCase    :							List[str]									=					self.module_class(config=__UpperCamelCase      ,		dtype=__UpperCamelCase      ,		**__UpperCamelCase     )
				super().__init__(__UpperCamelCase      ,		__UpperCamelCase      ,		input_shape=__UpperCamelCase      ,		seed=__UpperCamelCase      ,		dtype=__UpperCamelCase     )
			def 							__snake_case (		self      ,		_A      ,		_A      ,		_A = None     )    ->    str:






				'''simple docstring'''







				_UpperCAmelCase    :							Dict									=					jnp.zeros(input_shape[0]      ,		dtype="""i4"""     )
				_UpperCAmelCase    :							List[Any]									=					jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase     ).shape[-1]     )      ,		input_shape[0]     )
				_UpperCAmelCase    :							str									=					jnp.ones_like(__UpperCamelCase     )
				_UpperCAmelCase    :							Tuple									=					jnp.ones_like(__UpperCamelCase     )

				_UpperCAmelCase    :							Any									=					jax.random.normal(__UpperCamelCase      ,		input_shape[1]     )

				_UpperCAmelCase						,   _UpperCAmelCase    :							Tuple									=					jax.random.split(__UpperCamelCase     )
				_UpperCAmelCase    :							Tuple									=					{"""params""": params_rng, """dropout""": dropout_rng}

				return self.module.init(__UpperCamelCase      ,		__UpperCamelCase      ,		__UpperCamelCase      ,		__UpperCamelCase      ,		__UpperCamelCase      ,		__UpperCamelCase     )["params"]
			def __call__(		self      ,		_A      ,		_A      ,		_A=None      ,		_A=None      ,		_A=None      ,		_A = None      ,		_A = None      ,		_A = False      ,		_A = None      ,		_A = None      ,		_A = None      ,		)    ->    str:






				'''simple docstring'''







				_UpperCAmelCase    :							Any									=					output_attentions if output_attentions is not None else self.config.output_attentions
				_UpperCAmelCase    :							Dict									=					(
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				_UpperCAmelCase    :							Dict									=					return_dict if return_dict is not None else self.config.return_dict

				if position_ids is None:
					_UpperCAmelCase    :							Tuple									=					jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase     ).shape[-1]     )      ,		input_ids.shape     )

				if token_type_ids is None:
					_UpperCAmelCase    :							List[str]									=					jnp.zeros_like(__UpperCamelCase     )

				if attention_mask is None:
					_UpperCAmelCase    :							List[str]									=					jnp.ones_like(__UpperCamelCase     )

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

				return self.module.apply(
				    {"""params""": params or self.params}      ,		jnp.array(__UpperCamelCase      ,		dtype="""i4"""     )      ,		jnp.array(__UpperCamelCase      ,		dtype=jnp.floataa     )      ,		jnp.array(__UpperCamelCase      ,		dtype="""i4"""     )      ,		jnp.array(__UpperCamelCase      ,		dtype="""i4"""     )      ,		jnp.array(__UpperCamelCase      ,		dtype="""i4"""     )      ,		not train      ,		__UpperCamelCase      ,		__UpperCamelCase      ,		__UpperCamelCase      ,		rngs=__UpperCamelCase      ,		)
			def 							__snake_case (		self      ,		_A      ,		_A=None      ,		_A=None      ,		_A=None      ,		_A = None      ,		_A = None      ,		_A=False      ,		)    ->    str:






				'''simple docstring'''







				if position_ids is None:
					_UpperCAmelCase    :							List[str]									=					jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase     ).shape[-1]     )      ,		input_ids.shape     )

				if token_type_ids is None:
					_UpperCAmelCase    :							Any									=					jnp.zeros_like(__UpperCamelCase     )

				if attention_mask is None:
					_UpperCAmelCase    :							Optional[int]									=					jnp.ones_like(__UpperCamelCase     )

				# Handle any PRNG if needed
				_UpperCAmelCase    :							Any									=					{}
				if dropout_rng is not None:
					_UpperCAmelCase    :							str									=					dropout_rng

				def _get_features(_A      ,		_A      ,		_A      ,		_A      ,		_A      ,		_A     ):
					_UpperCAmelCase    :							int									=					module.text_model(
					    input_ids=__UpperCamelCase      ,		attention_mask=__UpperCamelCase      ,		position_ids=__UpperCamelCase      ,		token_type_ids=__UpperCamelCase      ,		deterministic=__UpperCamelCase      ,		)
					_UpperCAmelCase    :							str									=					text_outputs[1]
					_UpperCAmelCase    :							Dict									=					module.text_projection(__UpperCamelCase     )
					return text_features

				return self.module.apply(
				    {"""params""": params or self.params}      ,		jnp.array(__UpperCamelCase      ,		dtype="""i4"""     )      ,		jnp.array(__UpperCamelCase      ,		dtype="""i4"""     )      ,		jnp.array(__UpperCamelCase      ,		dtype="""i4"""     )      ,		jnp.array(__UpperCamelCase      ,		dtype="""i4"""     )      ,		not train      ,		method=_get_features      ,		rngs=__UpperCamelCase      ,		)
			def 							__snake_case (		self      ,		_A      ,		_A = None      ,		_A = None      ,		_A=False     )    ->    str:






				'''simple docstring'''







				_UpperCAmelCase    :							Any									=					{}
				if dropout_rng is not None:
					_UpperCAmelCase    :							Any									=					dropout_rng

				def _get_features(_A      ,		_A      ,		_A     ):
					_UpperCAmelCase    :							List[str]									=					module.vision_model(pixel_values=__UpperCamelCase      ,		deterministic=__UpperCamelCase     )
					_UpperCAmelCase    :							Optional[Any]									=					vision_outputs[1]  # pooled_output
					_UpperCAmelCase    :							str									=					module.visual_projection(__UpperCamelCase     )
					return image_features

				return self.module.apply(
				    {"""params""": params or self.params}      ,		jnp.array(__UpperCamelCase      ,		dtype=jnp.floataa     )      ,		not train      ,		method=_get_features      ,		rngs=__UpperCamelCase      ,		)
			@classmethod
			def 							__snake_case (		cls      ,		_A = None      ,		_A = None      ,		*_A      ,		**_A      ,		)    ->    Tuple:






				'''simple docstring'''







				_UpperCAmelCase    :							Tuple									=					{
				    argument[len("""text_"""     ) :]: value for argument, value in kwargs.items() if argument.startswith("""text_"""     )
				}

				_UpperCAmelCase    :							Dict									=					{
				    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
				_UpperCAmelCase    :							List[Any]									=					kwargs_text.pop("""model"""      ,		__UpperCamelCase     )
				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

						_UpperCAmelCase    :							Optional[int]									=					AutoConfig.from_pretrained(__UpperCamelCase     )
						_UpperCAmelCase    :							Optional[int]									=					text_config

					_UpperCAmelCase    :							Dict									=					FlaxAutoModel.from_pretrained(__UpperCamelCase      ,		*__UpperCamelCase      ,		**__UpperCamelCase     )

				_UpperCAmelCase    :							Dict									=					kwargs_vision.pop("""model"""      ,		__UpperCamelCase     )
				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

						_UpperCAmelCase    :							List[Any]									=					AutoConfig.from_pretrained(__UpperCamelCase     )
						_UpperCAmelCase    :							Optional[Any]									=					vision_config

					_UpperCAmelCase    :							List[Any]									=					FlaxAutoModel.from_pretrained(__UpperCamelCase      ,		*__UpperCamelCase      ,		**__UpperCamelCase     )

				# instantiate config with corresponding kwargs
				_UpperCAmelCase    :							Optional[Any]									=					kwargs.pop("""dtype"""      ,		jnp.floataa     )
				_UpperCAmelCase    :							List[str]									=					HybridCLIPConfig.from_text_vision_configs(text_model.config      ,		vision_model.config      ,		**__UpperCamelCase     )

				# init model
				_UpperCAmelCase    :							Optional[Any]									=					cls(__UpperCamelCase      ,		*__UpperCamelCase      ,		dtype=__UpperCamelCase      ,		**__UpperCamelCase     )

				if vision_config.model_type == "clip":
					_UpperCAmelCase    :							Optional[Any]									=					vision_model.params["""vision_model"""]
					_UpperCAmelCase    :							List[str]									=					vision_model.params["""visual_projection"""]["""kernel"""]
				else:
					_UpperCAmelCase    :							List[Any]									=					vision_model.params

				_UpperCAmelCase    :							Optional[Any]									=					text_model.params

				return model



"""simple docstring"""






from __future__ import annotations
def       UpperCamelCase					(			_lowerCAmelCase     :						list,  _lowerCAmelCase     :						list )				->			list:
	_UpperCAmelCase    :							List[Any]									=					[None] * (len(_lowerCAmelCase ) + len(_lowerCAmelCase ))

	_UpperCAmelCase    :							Union[str, Any]									=					0  # pointer to current index for left Half
	_UpperCAmelCase    :							str									=					0  # pointer to current index for the right Half
	_UpperCAmelCase    :							Optional[int]									=					0  # pointer to current index for the sorted array Half

	while pointera < len(_lowerCAmelCase ) and pointera < len(_lowerCAmelCase ):
		if left_half[pointera] < right_half[pointera]:
			_UpperCAmelCase    :							Tuple									=					left_half[pointera]
			pointera += 1
			index += 1
		else:
			_UpperCAmelCase    :							int									=					right_half[pointera]
			pointera += 1
			index += 1
	while pointera < len(_lowerCAmelCase ):
		_UpperCAmelCase    :							Tuple									=					left_half[pointera]
		pointera += 1
		index += 1

	while pointera < len(_lowerCAmelCase ):
		_UpperCAmelCase    :							Union[str, Any]									=					right_half[pointera]
		pointera += 1
		index += 1

	return sorted_array
def       UpperCamelCase					(			_lowerCAmelCase     :						list )				->			list:
	if len(_lowerCAmelCase ) <= 1:
		return array
	# the actual formula to calculate the middle element = left + (right - left) // 2
	# this avoids integer overflow in case of large N
	_UpperCAmelCase    :							Any									=					0 + (len(_lowerCAmelCase ) - 0) // 2

	# Split the array into halves till the array length becomes equal to One
	# merge the arrays of single length returned by mergeSort function and
	# pass them into the merge arrays function which merges the array
	_UpperCAmelCase    :							Optional[Any]									=					array[:middle]
	_UpperCAmelCase    :							Optional[int]									=					array[middle:]

	return merge(merge_sort(_lowerCAmelCase ),  merge_sort(_lowerCAmelCase ) )


if __name__ == "__main__":
						import doctest

						doctest.testmod()





