



import random
from typing import List, Optional, Union

import numpy as np
import torch

from ..state import AcceleratorState
from .constants import CUDA_DISTRIBUTED_TYPES
from .dataclasses import DistributedType, RNGType
from .imports import is_npu_available, is_tpu_available, is_xpu_available


if is_tpu_available(check_device=False):
      import torch_xla.core.xla_model as xm







def       lowerCamelCase_  (    _a       :						int   ,		_a       :						bool = False    ):



    '''simple docstring'''




    if device_specific:
        seed += AcceleratorState().process_index
    random.seed(_a    )
    np.random.seed(_a    )
    torch.manual_seed(_a    )
    if is_xpu_available():
        torch.xpu.manual_seed_all(_a    )
    elif is_npu_available():
        torch.npu.manual_seed_all(_a    )
    else:
        torch.cuda.manual_seed_all(_a    )
    # ^^ safe to call this function even if cuda is not available
    if is_tpu_available():
        xm.set_rng_state(_a    )







def       lowerCamelCase_  (    _a       :						Optional[RNGType] = None   ,		_a       :						Optional[torch.Generator] = None    ):



    '''simple docstring'''




    if rng_type == RNGType.TORCH:
        UpperCAmelCase_     :       Dict  		=					torch.get_rng_state()
    elif rng_type == RNGType.CUDA:
        UpperCAmelCase_     :       Any  		=					torch.cuda.get_rng_state()
    elif rng_type == RNGType.XLA:
        assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
        UpperCAmelCase_     :       List[str]  		=					torch.tensor(xm.get_rng_state()    )
    elif rng_type == RNGType.NPU:
        assert is_npu_available(), "Can't synchronize NPU seeds on an environment without NPUs."
        UpperCAmelCase_     :       List[Any]  		=					torch.npu.get_rng_state()
    elif rng_type == RNGType.XPU:
        assert is_xpu_available(), "Can't synchronize XPU seeds on an environment without XPUs."
        UpperCAmelCase_     :       int  		=					torch.xpu.get_rng_state()
    elif rng_type == RNGType.GENERATOR:
        assert generator is not None, "Need a generator to synchronize its seed."
        UpperCAmelCase_     :       List[str]  		=					generator.get_state()

    # Broadcast the rng state from device 0 to other devices
    UpperCAmelCase_     :       Any  		=					AcceleratorState()
    if state.distributed_type == DistributedType.TPU:
        UpperCAmelCase_     :       Tuple  		=					rng_state.to(xm.xla_device()    )
        xm.collective_broadcast([rng_state]    )
        xm.mark_step()
        UpperCAmelCase_     :       Tuple  		=					rng_state.cpu()
    elif (
        state.distributed_type in CUDA_DISTRIBUTED_TYPES
        or state.distributed_type == DistributedType.MULTI_NPU
        or state.distributed_type == DistributedType.MULTI_XPU
    ):
        UpperCAmelCase_     :       str  		=					rng_state.to(state.device    )
        torch.distributed.broadcast(_a   ,		0    )
        UpperCAmelCase_     :       Union[str, Any]  		=					rng_state.cpu()
    elif state.distributed_type == DistributedType.MULTI_CPU:
        torch.distributed.broadcast(_a   ,		0    )

    # Set the broadcast rng state
    if rng_type == RNGType.TORCH:
        torch.set_rng_state(_a    )
    elif rng_type == RNGType.CUDA:
        torch.cuda.set_rng_state(_a    )
    elif rng_type == RNGType.NPU:
        torch.npu.set_rng_state(_a    )
    elif rng_type == RNGType.XPU:
        torch.xpu.set_rng_state(_a    )
    elif rng_type == RNGType.XLA:
        xm.set_rng_state(rng_state.item()    )
    elif rng_type == RNGType.GENERATOR:
        generator.set_state(_a    )







def       lowerCamelCase_  (    _a       :						List[Union[str, RNGType]]   ,		_a       :						Optional[torch.Generator] = None    ):



    '''simple docstring'''




    for rng_type in rng_types:
        synchronize_rng_state(RNGType(_a    )   ,		generator=_a    )






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_      :							List[str]    							=					logging.get_logger(__name__)






class        __lowerCAmelCase					(			nn.Module					):
			snake_case_ :					HybridCLIPConfig
			snake_case_ :					jnp.dtype						=	jnp.floataa







			def UpperCamelCase     (			self       :     Union[str, Any] ):




									"""simple docstring"""


									_UpperCAmelCase						=	self.config.text_config
									_UpperCAmelCase						=	self.config.vision_config

									_UpperCAmelCase						=	self.config.projection_dim
									_UpperCAmelCase						=	text_config.hidden_size
									_UpperCAmelCase						=	vision_config.hidden_size

									_UpperCAmelCase						=	FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
									_UpperCAmelCase						=	FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__       ,       snake_case__ ).module_class

									_UpperCAmelCase						=	text_module(snake_case__       ,       dtype=self.dtype )
									_UpperCAmelCase						=	vision_module(snake_case__       ,       dtype=self.dtype )

									_UpperCAmelCase						=	nn.Dense(
									    self.projection_dim       ,       dtype=self.dtype       ,       kernel_init=jax.nn.initializers.normal(0.02 )       ,       use_bias=snake_case__       ,       )
									_UpperCAmelCase						=	nn.Dense(
									    self.projection_dim       ,       dtype=self.dtype       ,       kernel_init=jax.nn.initializers.normal(0.02 )       ,       use_bias=snake_case__       ,       )
									_UpperCAmelCase						=	self.param("logit_scale"       ,       jax.nn.initializers.ones       ,       [] )







			def __call__(			self       :     List[Any]       ,       snake_case__       :     List[Any]=None       ,       snake_case__       :     List[str]=None       ,       snake_case__       :     Optional[Any]=None       ,       snake_case__       :     Dict=None       ,       snake_case__       :     Tuple=None       ,       snake_case__       :     bool = True       ,       snake_case__       :     Dict=None       ,       snake_case__       :     Dict=None       ,       snake_case__       :     Dict=None       ,       ):




									"""simple docstring"""


									_UpperCAmelCase						=	return_dict if return_dict is not None else self.config.return_dict

									_UpperCAmelCase						=	self.vision_model(
									    pixel_values=snake_case__       ,       deterministic=snake_case__       ,       output_attentions=snake_case__       ,       output_hidden_states=snake_case__       ,       return_dict=snake_case__       ,       )

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

									_UpperCAmelCase						=	vision_outputs[1]
									_UpperCAmelCase						=	self.visual_projection(snake_case__ )

									_UpperCAmelCase						=	text_outputs[1]
									_UpperCAmelCase						=	self.text_projection(snake_case__ )

									# normalized features
									_UpperCAmelCase						=	image_embeds / jnp.linalg.norm(snake_case__       ,       axis=-1       ,       keepdims=snake_case__ )
									_UpperCAmelCase						=	text_embeds / jnp.linalg.norm(snake_case__       ,       axis=-1       ,       keepdims=snake_case__ )

									# cosine similarity as logits
									_UpperCAmelCase						=	jnp.exp(self.logit_scale )
									_UpperCAmelCase						=	jnp.matmul(snake_case__       ,       image_embeds.T ) * logit_scale
									_UpperCAmelCase						=	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=snake_case__       ,       logits_per_text=snake_case__       ,       text_embeds=snake_case__       ,       image_embeds=snake_case__       ,       text_model_output=snake_case__       ,       vision_model_output=snake_case__       ,       )






class        __lowerCAmelCase					(			UpperCAmelCase__					):
			snake_case_ :					Union[str, Any]						=	HybridCLIPConfig
			snake_case_ :					Optional[Any]						=	FlaxHybridCLIPModule







			def __init__(			self       :     List[Any]       ,       snake_case__       :     HybridCLIPConfig       ,       snake_case__       :     Optional[Tuple] = None       ,       snake_case__       :     int = 0       ,       snake_case__       :     jnp.dtype = jnp.floataa       ,       **snake_case__       :     Any       ,       ):




									"""simple docstring"""


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

									_UpperCAmelCase						=	self.module_class(config=snake_case__       ,       dtype=snake_case__       ,       **snake_case__ )
									super().__init__(snake_case__       ,       snake_case__       ,       input_shape=snake_case__       ,       seed=snake_case__       ,       dtype=snake_case__ )







			def UpperCamelCase     (			self       :     Union[str, Any]       ,       snake_case__       :     jax.random.PRNGKey       ,       snake_case__       :     Tuple       ,       snake_case__       :     FrozenDict = None ):




									"""simple docstring"""


									_UpperCAmelCase						=	jnp.zeros(input_shape[0]       ,       dtype="i4" )
									_UpperCAmelCase						=	jnp.broadcast_to(jnp.arange(jnp.atleast_ad(snake_case__ ).shape[-1] )       ,       input_shape[0] )
									_UpperCAmelCase						=	jnp.ones_like(snake_case__ )
									_UpperCAmelCase						=	jnp.ones_like(snake_case__ )

									_UpperCAmelCase						=	jax.random.normal(snake_case__       ,       input_shape[1] )

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

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







			def __call__(			self       :     Union[str, Any]       ,       snake_case__       :     Optional[int]       ,       snake_case__       :     Tuple       ,       snake_case__       :     Optional[int]=None       ,       snake_case__       :     Tuple=None       ,       snake_case__       :     Optional[Any]=None       ,       snake_case__       :     dict = None       ,       snake_case__       :     jax.random.PRNGKey = None       ,       snake_case__       :     bool = False       ,       snake_case__       :     Optional[bool] = None       ,       snake_case__       :     Optional[bool] = None       ,       snake_case__       :     Optional[bool] = None       ,       ):




									"""simple docstring"""


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

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

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

									if attention_mask is None:
															_UpperCAmelCase						=	jnp.ones_like(snake_case__ )

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

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







			def UpperCamelCase     (			self       :     int       ,       snake_case__       :     int       ,       snake_case__       :     List[str]=None       ,       snake_case__       :     Union[str, Any]=None       ,       snake_case__       :     str=None       ,       snake_case__       :     dict = None       ,       snake_case__       :     jax.random.PRNGKey = None       ,       snake_case__       :     int=False       ,       ):




									"""simple docstring"""


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

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

									if attention_mask is None:
															_UpperCAmelCase						=	jnp.ones_like(snake_case__ )

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

									def _get_features(snake_case__       :     Union[str, Any]       ,       snake_case__       :     Optional[Any]       ,       snake_case__       :     List[Any]       ,       snake_case__       :     int       ,       snake_case__       :     Dict       ,       snake_case__       :     List[str] ):
															_UpperCAmelCase						=	module.text_model(
															    input_ids=snake_case__       ,       attention_mask=snake_case__       ,       position_ids=snake_case__       ,       token_type_ids=snake_case__       ,       deterministic=snake_case__       ,       )
															_UpperCAmelCase						=	text_outputs[1]
															_UpperCAmelCase						=	module.text_projection(snake_case__ )
															return text_features

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







			def UpperCamelCase     (			self       :     List[Any]       ,       snake_case__       :     Optional[Any]       ,       snake_case__       :     dict = None       ,       snake_case__       :     jax.random.PRNGKey = None       ,       snake_case__       :     str=False ):




									"""simple docstring"""


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

									def _get_features(snake_case__       :     Optional[Any]       ,       snake_case__       :     Tuple       ,       snake_case__       :     int ):
															_UpperCAmelCase						=	module.vision_model(pixel_values=snake_case__       ,       deterministic=snake_case__ )
															_UpperCAmelCase						=	vision_outputs[1]  # pooled_output
															_UpperCAmelCase						=	module.visual_projection(snake_case__ )
															return image_features

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







			@classmethod
			def UpperCamelCase     (			cls       :     Tuple       ,       snake_case__       :     str = None       ,       snake_case__       :     str = None       ,       *snake_case__       :     Any       ,       **snake_case__       :     Tuple       ,       ):




									"""simple docstring"""


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

									_UpperCAmelCase						=	{
									    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						=	kwargs_text.pop("model"       ,       snake_case__ )
									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						=	AutoConfig.from_pretrained(snake_case__ )
																					_UpperCAmelCase						=	text_config

															_UpperCAmelCase						=	FlaxAutoModel.from_pretrained(snake_case__       ,       *snake_case__       ,       **snake_case__ )

									_UpperCAmelCase						=	kwargs_vision.pop("model"       ,       snake_case__ )
									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						=	AutoConfig.from_pretrained(snake_case__ )
																					_UpperCAmelCase						=	vision_config

															_UpperCAmelCase						=	FlaxAutoModel.from_pretrained(snake_case__       ,       *snake_case__       ,       **snake_case__ )

									# instantiate config with corresponding kwargs
									_UpperCAmelCase						=	kwargs.pop("dtype"       ,       jnp.floataa )
									_UpperCAmelCase						=	HybridCLIPConfig.from_text_vision_configs(text_model.config       ,       vision_model.config       ,       **snake_case__ )

									# init model
									_UpperCAmelCase						=	cls(snake_case__       ,       *snake_case__       ,       dtype=snake_case__       ,       **snake_case__ )

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

									_UpperCAmelCase						=	text_model.params

									return model




