code: 
  
  
  
  
  '''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   lowerCamelCase_							( nn.Module							):
  
  
  
  
  					'''simple docstring'''
  
  
  
  
  
  
  					lowercase__					=				4_2
  					lowercase__					=				jnp.floataa
  
  
  
  
  
  
  					def        __SCREAMING_SNAKE_CASE			(					self				):
  							_snake_case			 = self.config.text_config
  							_snake_case			 = self.config.vision_config
  
  							_snake_case			 = self.config.projection_dim
  							_snake_case			 = text_config.hidden_size
  							_snake_case			 = vision_config.hidden_size
  
  							_snake_case			 = FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
  							_snake_case			 = FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__,						__UpperCamelCase				).module_class
  
  							_snake_case			 = text_module(__UpperCamelCase,						dtype=self.dtype				)
  							_snake_case			 = vision_module(__UpperCamelCase,						dtype=self.dtype				)
  
  							_snake_case			 = nn.Dense(
  							    self.projection_dim,						dtype=self.dtype,						kernel_init=jax.nn.initializers.normal(0.02				),						use_bias=__UpperCamelCase,						)
  							_snake_case			 = nn.Dense(
  							    self.projection_dim,						dtype=self.dtype,						kernel_init=jax.nn.initializers.normal(0.02				),						use_bias=__UpperCamelCase,						)
  							_snake_case			 = self.param("logit_scale",						jax.nn.initializers.ones,						[]				)
  
  
  
  
  
  
  					def __call__(					self,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase = True,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						):
  							_snake_case			 = return_dict if return_dict is not None else self.config.return_dict
  
  							_snake_case			 = self.vision_model(
  							    pixel_values=__UpperCamelCase,						deterministic=__UpperCamelCase,						output_attentions=__UpperCamelCase,						output_hidden_states=__UpperCamelCase,						return_dict=__UpperCamelCase,						)
  
  							_snake_case			 = 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,						)
  
  							_snake_case			 = vision_outputs[1]
  							_snake_case			 = self.visual_projection(__UpperCamelCase				)
  
  							_snake_case			 = text_outputs[1]
  							_snake_case			 = self.text_projection(__UpperCamelCase				)
  
  							# normalized features
  							_snake_case			 = image_embeds / jnp.linalg.norm(__UpperCamelCase,						axis=-1,						keepdims=__UpperCamelCase				)
  							_snake_case			 = text_embeds / jnp.linalg.norm(__UpperCamelCase,						axis=-1,						keepdims=__UpperCamelCase				)
  
  							# cosine similarity as logits
  							_snake_case			 = jnp.exp(self.logit_scale				)
  							_snake_case			 = jnp.matmul(__UpperCamelCase,						image_embeds.T				) * logit_scale
  							_snake_case			 = 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   lowerCamelCase_							( UpperCAmelCase_							):
  
  
  
  
  					'''simple docstring'''
  
  
  
  
  
  
  					lowercase__					=				HybridCLIPConfig
  					lowercase__					=				FlaxHybridCLIPModule
  
  
  
  
  
  
  					def __init__(					self,						__UpperCamelCase,						__UpperCamelCase = None,						__UpperCamelCase = 0,						__UpperCamelCase = jnp.floataa,						**__UpperCamelCase,						):
  							if input_shape is None:
  									_snake_case			 = ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))
  
  							_snake_case			 = self.module_class(config=__UpperCamelCase,						dtype=__UpperCamelCase,						**__UpperCamelCase				)
  							super().__init__(__UpperCamelCase,						__UpperCamelCase,						input_shape=__UpperCamelCase,						seed=__UpperCamelCase,						dtype=__UpperCamelCase				)
  
  
  
  
  
  
  					def        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase = None				):
  							# init input tensor
  							_snake_case			 = jnp.zeros(input_shape[0],						dtype="i4"				)
  							_snake_case			 = jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase				).shape[-1]				),						input_shape[0]				)
  							_snake_case			 = jnp.ones_like(__UpperCamelCase				)
  							_snake_case			 = jnp.ones_like(__UpperCamelCase				)
  
  							_snake_case			 = jax.random.normal(__UpperCamelCase,						input_shape[1]				)
  
  							_snake_case ,					_snake_case			 = jax.random.split(__UpperCamelCase				)
  							_snake_case			 = {"params": params_rng, "dropout": dropout_rng}
  
  							return self.module.init(__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase				)["params"]
  
  
  
  
  
  
  					def __call__(					self,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase = None,						__UpperCamelCase = None,						__UpperCamelCase = False,						__UpperCamelCase = None,						__UpperCamelCase = None,						__UpperCamelCase = None,						):
  							_snake_case			 = output_attentions if output_attentions is not None else self.config.output_attentions
  							_snake_case			 = (
  							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
  							)
  							_snake_case			 = return_dict if return_dict is not None else self.config.return_dict
  
  							if position_ids is None:
  									_snake_case			 = jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase				).shape[-1]				),						input_ids.shape				)
  
  							if token_type_ids is None:
  									_snake_case			 = jnp.zeros_like(__UpperCamelCase				)
  
  							if attention_mask is None:
  									_snake_case			 = jnp.ones_like(__UpperCamelCase				)
  
  							# Handle any PRNG if needed
  							_snake_case			 = {}
  							if dropout_rng is not None:
  									_snake_case			 = 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        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase=None,						__UpperCamelCase = None,						__UpperCamelCase = None,						__UpperCamelCase=False,						):
  							if position_ids is None:
  									_snake_case			 = jnp.broadcast_to(jnp.arange(jnp.atleast_ad(__UpperCamelCase				).shape[-1]				),						input_ids.shape				)
  
  							if token_type_ids is None:
  									_snake_case			 = jnp.zeros_like(__UpperCamelCase				)
  
  							if attention_mask is None:
  									_snake_case			 = jnp.ones_like(__UpperCamelCase				)
  
  							# Handle any PRNG if needed
  							_snake_case			 = {}
  							if dropout_rng is not None:
  									_snake_case			 = dropout_rng
  
  							def _get_features(__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase				):
  									_snake_case			 = module.text_model(
  									    input_ids=__UpperCamelCase,						attention_mask=__UpperCamelCase,						position_ids=__UpperCamelCase,						token_type_ids=__UpperCamelCase,						deterministic=__UpperCamelCase,						)
  									_snake_case			 = text_outputs[1]
  									_snake_case			 = 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        __SCREAMING_SNAKE_CASE			(					self,						__UpperCamelCase,						__UpperCamelCase = None,						__UpperCamelCase = None,						__UpperCamelCase=False				):
  							_snake_case			 = {}
  							if dropout_rng is not None:
  									_snake_case			 = dropout_rng
  
  							def _get_features(__UpperCamelCase,						__UpperCamelCase,						__UpperCamelCase				):
  									_snake_case			 = module.vision_model(pixel_values=__UpperCamelCase,						deterministic=__UpperCamelCase				)
  									_snake_case			 = vision_outputs[1]  # pooled_output
  									_snake_case			 = 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        __SCREAMING_SNAKE_CASE			(					cls,						__UpperCamelCase = None,						__UpperCamelCase = None,						*__UpperCamelCase,						**__UpperCamelCase,						):
  							_snake_case			 = {
  							    argument[len("text_"				) :]: value for argument, value in kwargs.items() if argument.startswith("text_"				)
  							}
  
  							_snake_case			 = {
  							    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
  							_snake_case			 = 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
  
  											_snake_case			 = AutoConfig.from_pretrained(__UpperCamelCase				)
  											_snake_case			 = text_config
  
  									_snake_case			 = FlaxAutoModel.from_pretrained(__UpperCamelCase,						*__UpperCamelCase,						**__UpperCamelCase				)
  
  							_snake_case			 = 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
  
  											_snake_case			 = AutoConfig.from_pretrained(__UpperCamelCase				)
  											_snake_case			 = vision_config
  
  									_snake_case			 = FlaxAutoModel.from_pretrained(__UpperCamelCase,						*__UpperCamelCase,						**__UpperCamelCase				)
  
  							# instantiate config with corresponding kwargs
  							_snake_case			 = kwargs.pop("dtype",						jnp.floataa				)
  							_snake_case			 = HybridCLIPConfig.from_text_vision_configs(text_model.config,						vision_model.config,						**__UpperCamelCase				)
  
  							# init model
  							_snake_case			 = cls(__UpperCamelCase,						*__UpperCamelCase,						dtype=__UpperCamelCase,						**__UpperCamelCase				)
  
  							if vision_config.model_type == "clip":
  									_snake_case			 = vision_model.params["vision_model"]
  									_snake_case			 = vision_model.params["visual_projection"]["kernel"]
  							else:
  									_snake_case			 = vision_model.params
  
  							_snake_case			 = text_model.params
  
  							return model
  
  
  
  
  
  
style_context: 
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  
  
  import importlib
  import importlib.metadata
  import os
  import warnings
  from distutils.util import strtobool
  from functools import lru_cache
  
  import torch
  from packaging import version
  from packaging.version import parse
  
  from .environment import parse_flag_from_env
  from .versions import compare_versions, is_torch_version
  
  
  try:
  						import torch_xla.core.xla_model as xm  # noqa: F401
  
  						_lowerCamelCase							:		int										=						True
  except ImportError:
  						_lowerCamelCase							:		List[str]										=						False
  
  
  # Cache this result has it's a C FFI call which can be pretty time-consuming
  _lowerCamelCase							:		List[Any]										=						torch.distributed.is_available()
  
  
  
  
  
  
  def 	_lowercase	(			lowercase_						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = importlib.util.find_spec(lowercase_						) is not None
  		if package_exists:
  				try:
  						_snake_case			 = importlib.metadata.metadata(lowercase_						)
  						return True
  				except importlib.metadata.PackageNotFoundError:
  						return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _torch_distributed_available
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		try:
  				pass
  		except ImportError:
  				print(
  				    "Intel(R) oneCCL Bindings for PyTorch* is required to run DDP on Intel(R) GPUs, but it is not"
  				    " detected. If you see \"ValueError: Invalid backend: 'ccl'\" error, please install Intel(R) oneCCL"
  				    " Bindings for PyTorch*."						)
  		return (
  		    importlib.util.find_spec("torch_ccl"						) is not None
  		    or importlib.util.find_spec("oneccl_bindings_for_pytorch"						) is not None
  		)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return importlib.metadata.version("oneccl_bind_pt"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("transformer_engine"						)
  
  
  
  
  
  
  @lru_cache
  def 	_lowercase	(			lowercase_=True						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if torch.cuda.is_available():
  				return False
  		if _tpu_available and check_device:
  				try:
  						# Will raise a RuntimeError if no XLA configuration is found
  						_snake_case			 = xm.xla_device()
  						return True
  				except RuntimeError:
  						return False
  		return _tpu_available
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("deepspeed"						)
  
  
  
  
  
  
  def 	_lowercase	(			lowercase_=False						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if is_tpu_available():
  				return not ignore_tpu
  		if torch.cuda.is_available():
  				return torch.cuda.is_bfaa_supported()
  		if is_npu_available():
  				return False
  		return True
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = _is_package_available("bitsandbytes"						)
  		if package_exists:
  				_snake_case			 = version.parse(importlib.metadata.version("bitsandbytes"						)						)
  				return compare_versions(lowercase_		,							">="		,							"0.39.0"						)
  		return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = _is_package_available("bitsandbytes"						)
  		if package_exists:
  				_snake_case			 = version.parse(importlib.metadata.version("bitsandbytes"						)						)
  				return compare_versions(lowercase_		,							">="		,							"0.37.2"						)
  		return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("bitsandbytes"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if strtobool(os.environ.get("ACCELERATE_USE_MEGATRON_LM"		,							"False"						)						) == 1:
  				_snake_case			 = importlib.util.find_spec("megatron"						) is not None
  				if package_exists:
  						try:
  								_snake_case			 = parse(importlib.metadata.version("megatron-lm"						)						)
  								return compare_versions(lowercase_		,							">="		,							"2.2.0"						)
  						except Exception as e:
  								warnings.warn(f'''Parse Megatron version failed. Exception:{e}'''						)
  								return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("safetensors"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("transformers"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("datasets"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("aim"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("tensorboard"						) or _is_package_available("tensorboardX"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("wandb"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("comet_ml"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("boto3"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if _is_package_available("rich"						):
  				if "ACCELERATE_DISABLE_RICH" in os.environ:
  						warnings.warn(
  						    "`ACCELERATE_DISABLE_RICH` is deprecated and will be removed in v0.22.0 and deactivated by default. Please use `ACCELERATE_ENABLE_RICH` if you wish to use `rich`."						)
  						return not parse_flag_from_env("ACCELERATE_DISABLE_RICH"		,							lowercase_						)
  				return parse_flag_from_env("ACCELERATE_ENABLE_RICH"		,							lowercase_						)
  		return False
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("sagemaker"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("tqdm"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return _is_package_available("mlflow"						)
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		return is_torch_version(">="		,							"1.12"						) and torch.backends.mps.is_available() and torch.backends.mps.is_built()
  
  
  
  
  
  
  def 	_lowercase	(			):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		def get_major_and_minor_from_version(lowercase_						):
  				return str(version.parse(lowercase_						).major						) + "." + str(version.parse(lowercase_						).minor						)
  
  		_snake_case			 = importlib.metadata.version("torch"						)
  		if importlib.util.find_spec("intel_extension_for_pytorch"						) is None:
  				return False
  		_snake_case			 = "N/A"
  		try:
  				_snake_case			 = importlib.metadata.version("intel_extension_for_pytorch"						)
  		except importlib.metadata.PackageNotFoundError:
  				return False
  		_snake_case			 = get_major_and_minor_from_version(_torch_version						)
  		_snake_case			 = get_major_and_minor_from_version(_ipex_version						)
  		if torch_major_and_minor != ipex_major_and_minor:
  				warnings.warn(
  				    f'''Intel Extension for PyTorch {ipex_major_and_minor} needs to work with PyTorch {ipex_major_and_minor}.*,'''
  				    f''' but PyTorch {_torch_version} is found. Please switch to the matching version and run again.'''						)
  				return False
  		return True
  
  
  
  
  
  
  @lru_cache
  def 	_lowercase	(			lowercase_=False						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if importlib.util.find_spec("torch"						) is None or importlib.util.find_spec("torch_npu"						) is None:
  				return False
  
  		import torch
  		import torch_npu  # noqa: F401
  
  		if check_device:
  				try:
  						# Will raise a RuntimeError if no NPU is found
  						_snake_case			 = torch.npu.device_count()
  						return torch.npu.is_available()
  				except RuntimeError:
  						return False
  		return hasattr(lowercase_		,							"npu"						) and torch.npu.is_available()
  
  
  
  
  
  
  @lru_cache
  def 	_lowercase	(			lowercase_=False						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		if not parse_flag_from_env("ACCELERATE_USE_XPU"		,							default=lowercase_						):
  				return False
  		if is_ipex_available():
  				import torch
  
  				if is_torch_version("<="		,							"1.12"						):
  						return False
  		else:
  				return False
  
  		import intel_extension_for_pytorch  # noqa: F401
  
  		if check_device:
  				try:
  						# Will raise a RuntimeError if no XPU  is found
  						_snake_case			 = torch.xpu.device_count()
  						return torch.xpu.is_available()
  				except RuntimeError:
  						return False
  		return hasattr(lowercase_		,							"xpu"						) and torch.xpu.is_available()
  
  
  
  
  
  