





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


A__:  str       					=					logging.get_logger(__name__)




class    _a  (  nn.Module):

      """simple docstring"""


      UpperCamelCase__       						=   42
      UpperCamelCase__       						=   jnp.floataa






      def 	UpperCAmelCase_   (						self:	Dict			):
          '''simple docstring'''


          UpperCamelCase__:					List[str]												=      self.config.text_config
          UpperCamelCase__:					Any												=      self.config.vision_config

          UpperCamelCase__:					str												=      self.config.projection_dim
          UpperCamelCase__:					List[Any]												=      text_config.hidden_size
          UpperCamelCase__:					Dict												=      vision_config.hidden_size

          UpperCamelCase__:					Union[str, Any]												=      FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
          UpperCamelCase__:					str												=      FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__					,  _UpperCamelCase			).module_class

          UpperCamelCase__:					Union[str, Any]												=      text_module(_UpperCamelCase					,  dtype=self.dtype			)
          UpperCamelCase__:					Dict												=      vision_module(_UpperCamelCase					,  dtype=self.dtype			)

          UpperCamelCase__:					Any												=      nn.Dense(
              self.projection_dim					,  dtype=self.dtype					,  kernel_init=jax.nn.initializers.normal(0.02			)					,  use_bias=_UpperCamelCase					,  )
          UpperCamelCase__:					Union[str, Any]												=      nn.Dense(
              self.projection_dim					,  dtype=self.dtype					,  kernel_init=jax.nn.initializers.normal(0.02			)					,  use_bias=_UpperCamelCase					,  )
          UpperCamelCase__:					Optional[Any]												=      self.param("logit_scale"					,  jax.nn.initializers.ones					,  []			)






      def __call__(						self:	Optional[int]					,  __lowerCamelCase:	int=None					,  __lowerCamelCase:	int=None					,  __lowerCamelCase:	Optional[int]=None					,  __lowerCamelCase:	Optional[Any]=None					,  __lowerCamelCase:	Optional[int]=None					,  __lowerCamelCase:	bool = True					,  __lowerCamelCase:	str=None					,  __lowerCamelCase:	Any=None					,  __lowerCamelCase:	str=None					,  ):
          '''simple docstring'''


          UpperCamelCase__:					List[str]												=      return_dict if return_dict is not None else self.config.return_dict

          UpperCamelCase__:					int												=      self.vision_model(
              pixel_values=_UpperCamelCase					,  deterministic=_UpperCamelCase					,  output_attentions=_UpperCamelCase					,  output_hidden_states=_UpperCamelCase					,  return_dict=_UpperCamelCase					,  )

          UpperCamelCase__:					Optional[int]												=      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[int]												=      vision_outputs[1]
          UpperCamelCase__:					int												=      self.visual_projection(_UpperCamelCase			)

          UpperCamelCase__:					Optional[Any]												=      text_outputs[1]
          UpperCamelCase__:					Tuple												=      self.text_projection(_UpperCamelCase			)

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

          # cosine similarity as logits
          UpperCamelCase__:					Union[str, Any]												=      jnp.exp(self.logit_scale			)
          UpperCamelCase__:					Tuple												=      jnp.matmul(_UpperCamelCase					,  image_embeds.T			) * logit_scale
          UpperCamelCase__:					Optional[int]												=      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    _a  (  __lowerCamelCase):

      """simple docstring"""


      UpperCamelCase__       						=   HybridCLIPConfig
      UpperCamelCase__       						=   FlaxHybridCLIPModule






      def __init__(						self:	str					,  __lowerCamelCase:	HybridCLIPConfig					,  __lowerCamelCase:	Optional[Tuple] = None					,  __lowerCamelCase:	int = 0					,  __lowerCamelCase:	jnp.dtype = jnp.floataa					,  **__lowerCamelCase:	str					,  ):
          '''simple docstring'''


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

          UpperCamelCase__:					Optional[Any]												=      self.module_class(config=_UpperCamelCase					,  dtype=_UpperCamelCase					,  **_UpperCamelCase			)
          super().__init__(_UpperCamelCase					,  _UpperCamelCase					,  input_shape=_UpperCamelCase					,  seed=_UpperCamelCase					,  dtype=_UpperCamelCase			)






      def 	UpperCAmelCase_   (						self:	str					,  __lowerCamelCase:	jax.random.PRNGKey					,  __lowerCamelCase:	Tuple					,  __lowerCamelCase:	FrozenDict = None			):
          '''simple docstring'''


          UpperCamelCase__:					List[str]												=      jnp.zeros(input_shape[0]					,  dtype="i4"			)
          UpperCamelCase__:					List[str]												=      jnp.broadcast_to(jnp.arange(jnp.atleast_ad(_UpperCamelCase			).shape[-1]			)					,  input_shape[0]			)
          UpperCamelCase__:					List[Any]												=      jnp.ones_like(_UpperCamelCase			)
          UpperCamelCase__:					Optional[int]												=      jnp.ones_like(_UpperCamelCase			)

          UpperCamelCase__:					Union[str, Any]												=      jax.random.normal(_UpperCamelCase					,  input_shape[1]			)

          UpperCamelCase__							,				UpperCamelCase__:					int												=      jax.random.split(_UpperCamelCase			)
          UpperCamelCase__:					Union[str, Any]												=      {"params": params_rng, "dropout": dropout_rng}

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






      def __call__(						self:	Any					,  __lowerCamelCase:	str					,  __lowerCamelCase:	Any					,  __lowerCamelCase:	List[str]=None					,  __lowerCamelCase:	Any=None					,  __lowerCamelCase:	Union[str, Any]=None					,  __lowerCamelCase:	dict = None					,  __lowerCamelCase:	jax.random.PRNGKey = None					,  __lowerCamelCase:	bool = False					,  __lowerCamelCase:	Optional[bool] = None					,  __lowerCamelCase:	Optional[bool] = None					,  __lowerCamelCase:	Optional[bool] = None					,  ):
          '''simple docstring'''


          UpperCamelCase__:					List[str]												=      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__:					Tuple												=      return_dict if return_dict is not None else self.config.return_dict

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

          if token_type_ids is None:
              UpperCamelCase__:					int												=      jnp.zeros_like(_UpperCamelCase			)

          if attention_mask is None:
              UpperCamelCase__:					Union[str, Any]												=      jnp.ones_like(_UpperCamelCase			)

          # Handle any PRNG if needed
          UpperCamelCase__:					int												=      {}
          if dropout_rng is not None:
              UpperCamelCase__:					List[str]												=      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 	UpperCAmelCase_   (						self:	Tuple					,  __lowerCamelCase:	Optional[Any]					,  __lowerCamelCase:	Any=None					,  __lowerCamelCase:	Tuple=None					,  __lowerCamelCase:	Tuple=None					,  __lowerCamelCase:	dict = None					,  __lowerCamelCase:	jax.random.PRNGKey = None					,  __lowerCamelCase:	List[Any]=False					,  ):
          '''simple docstring'''


          if position_ids is None:
              UpperCamelCase__:					Union[str, Any]												=      jnp.broadcast_to(jnp.arange(jnp.atleast_ad(_UpperCamelCase			).shape[-1]			)					,  input_ids.shape			)

          if token_type_ids is None:
              UpperCamelCase__:					str												=      jnp.zeros_like(_UpperCamelCase			)

          if attention_mask is None:
              UpperCamelCase__:					str												=      jnp.ones_like(_UpperCamelCase			)

          # Handle any PRNG if needed
          UpperCamelCase__:					int												=      {}
          if dropout_rng is not None:
              UpperCamelCase__:					Dict												=      dropout_rng

          def _get_features(__lowerCamelCase:	List[str]					,  __lowerCamelCase:	List[Any]					,  __lowerCamelCase:	Any					,  __lowerCamelCase:	Dict					,  __lowerCamelCase:	Optional[Any]					,  __lowerCamelCase:	List[str]			):
              UpperCamelCase__:					int												=      module.text_model(
                  input_ids=_UpperCamelCase					,  attention_mask=_UpperCamelCase					,  position_ids=_UpperCamelCase					,  token_type_ids=_UpperCamelCase					,  deterministic=_UpperCamelCase					,  )
              UpperCamelCase__:					List[Any]												=      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 	UpperCAmelCase_   (						self:	Union[str, Any]					,  __lowerCamelCase:	List[Any]					,  __lowerCamelCase:	dict = None					,  __lowerCamelCase:	jax.random.PRNGKey = None					,  __lowerCamelCase:	Union[str, Any]=False			):
          '''simple docstring'''


          UpperCamelCase__:					Optional[Any]												=      {}
          if dropout_rng is not None:
              UpperCamelCase__:					int												=      dropout_rng

          def _get_features(__lowerCamelCase:	Optional[int]					,  __lowerCamelCase:	Union[str, Any]					,  __lowerCamelCase:	int			):
              UpperCamelCase__:					str												=      module.vision_model(pixel_values=_UpperCamelCase					,  deterministic=_UpperCamelCase			)
              UpperCamelCase__:					List[str]												=      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 	UpperCAmelCase_   (						cls:	List[str]					,  __lowerCamelCase:	str = None					,  __lowerCamelCase:	str = None					,  *__lowerCamelCase:	str					,  **__lowerCamelCase:	Optional[Any]					,  ):
          '''simple docstring'''


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

          UpperCamelCase__:					List[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
          UpperCamelCase__:					List[str]												=      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__:					Any												=      AutoConfig.from_pretrained(_UpperCamelCase			)
                  UpperCamelCase__:					Tuple												=      text_config

              UpperCamelCase__:					List[str]												=      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__:					Optional[int]												=      AutoConfig.from_pretrained(_UpperCamelCase			)
                  UpperCamelCase__:					Any												=      vision_config

              UpperCamelCase__:					int												=      FlaxAutoModel.from_pretrained(_UpperCamelCase					,  *_UpperCamelCase					,  **_UpperCamelCase			)

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

          # init model
          UpperCamelCase__:					List[Any]												=      cls(_UpperCamelCase					,  *_UpperCamelCase					,  dtype=_UpperCamelCase					,  **_UpperCamelCase			)

          if vision_config.model_type == "clip":
              UpperCamelCase__:					str												=      vision_model.params["vision_model"]
              UpperCamelCase__:					Union[str, Any]												=      vision_model.params["visual_projection"]["kernel"]
          else:
              UpperCamelCase__:					Tuple												=      vision_model.params

          UpperCamelCase__:					int												=      text_model.params

          return model








from .chunk_utils import chunk_layer
from .data_transforms import make_atomaa_masks
from .feats import atomaa_to_atomaa, frames_and_literature_positions_to_atomaa_pos, torsion_angles_to_frames
from .loss import compute_predicted_aligned_error, compute_tm
from .protein import Protein as OFProtein
from .protein import to_pdb
from .rigid_utils import Rigid, Rotation
from .tensor_utils import dict_multimap, flatten_final_dims, permute_final_dims

