




"""simple docstring"""



import copy
from typing import TYPE_CHECKING, Any, Mapping, Optional, OrderedDict

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig


if TYPE_CHECKING:
      from ... import PreTrainedTokenizerBase, TensorType

SCREAMING_SNAKE_CASE    =	logging.get_logger(__name__)




class      UpperCAmelCase_  (						lowerCAmelCase_							):
    lowercase__						 =				'vision-encoder-decoder'
    lowercase__						 =				True




    def __init__(   self						:	int     ,			**snake_case_						:	Any				)				->							List[str]:




     '''simple docstring'''

     super().__init__(**__lowerCAmelCase				)
     if "encoder" not in kwargs or "decoder" not in kwargs:
      raise ValueError(
          F"""A configuraton of type {self.model_type} cannot be instantiated because """
          F"""not both `encoder` and `decoder` sub-configurations are passed, but only {kwargs}"""				)

     A__	=    kwargs.pop("encoder"				)
     A__	=    encoder_config.pop("model_type"				)
     A__	=    kwargs.pop("decoder"				)
     A__	=    decoder_config.pop("model_type"				)

     A__	=    AutoConfig.for_model(__lowerCAmelCase     ,			**__lowerCAmelCase				)
     A__	=    AutoConfig.for_model(__lowerCAmelCase     ,			**__lowerCAmelCase				)
     A__	=    True




    @classmethod
    def 							__magic_name__							(   cls						:	int     ,			snake_case_						:	PretrainedConfig     ,			snake_case_						:	PretrainedConfig     ,			**snake_case_						:	str				)				->							Optional[int]:




     '''simple docstring'''

     logger.info("Setting `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config"				)
     A__	=    True
     A__	=    True

     return cls(encoder=encoder_config.to_dict()     ,			decoder=decoder_config.to_dict()     ,			**__lowerCAmelCase				)




    def 							__magic_name__							(   self						:	int				)				->							Any:




     '''simple docstring'''

     A__	=    copy.deepcopy(self.__dict__				)
     A__	=    self.encoder.to_dict()
     A__	=    self.decoder.to_dict()
     A__	=    self.__class__.model_type
     return output






class      UpperCAmelCase_  (						lowerCAmelCase_							):
    lowercase__						 =				version.parse('''1.11'''							)




    @property
    def 							__magic_name__							(   self						:	int				)				->							int:




     '''simple docstring'''

     return OrderedDict(
         [
             ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
         ]				)




    @property
    def 							__magic_name__							(   self						:	Tuple				)				->							int:




     '''simple docstring'''

     return 1e-4




    @property
    def 							__magic_name__							(   self						:	Dict				)				->							Tuple:




     '''simple docstring'''

     return OrderedDict({"last_hidden_state": {0: "batch", 1: "encoder_sequence"}}				)






class      UpperCAmelCase_  (						lowerCAmelCase_							):




    @property
    def 							__magic_name__							(   self						:	Any				)				->							Optional[Any]:




     '''simple docstring'''

     A__	=    OrderedDict()
     A__	=    {0: "batch", 1: "past_decoder_sequence + sequence"}
     A__	=    {0: "batch", 1: "past_decoder_sequence + sequence"}
     A__	=    {0: "batch", 1: "encoder_sequence"}

     return common_inputs




    def 							__magic_name__							(   self						:	Union[str, Any]     ,			snake_case_						:	"PreTrainedTokenizerBase"     ,			snake_case_						:	int = -1     ,			snake_case_						:	int = -1     ,			snake_case_						:	bool = False     ,			snake_case_						:	Optional["TensorType"] = None     ,			)				->							Any:




     '''simple docstring'''

     import torch

     A__	=    OrderedDict()

     A__	=    super().generate_dummy_inputs(
         __lowerCAmelCase     ,			batch_size=__lowerCAmelCase     ,			seq_length=__lowerCAmelCase     ,			is_pair=__lowerCAmelCase     ,			framework=__lowerCAmelCase				)

     A__,				A__	=    dummy_input["input_ids"].shape
     A__	=    (batch, encoder_sequence, self._config.encoder_hidden_size)
     A__	=    dummy_input.pop("input_ids"				)
     A__	=    dummy_input.pop("attention_mask"				)
     A__	=    torch.zeros(__lowerCAmelCase				)

     return common_inputs






class      UpperCAmelCase_  (						lowerCAmelCase_							):




    @property
    def 							__magic_name__							(   self						:	Tuple				)				->							Union[str, Any]:




     '''simple docstring'''

     pass




    def 							__magic_name__							(   self						:	List[Any]     ,			snake_case_						:	PretrainedConfig				)				->							List[Any]:




     '''simple docstring'''

     return VisionEncoderDecoderEncoderOnnxConfig(__lowerCAmelCase				)




    def 							__magic_name__							(   self						:	int     ,			snake_case_						:	PretrainedConfig     ,			snake_case_						:	PretrainedConfig     ,			snake_case_						:	str = "default"				)				->							int:




     '''simple docstring'''

     A__	=    encoder_config.hidden_size
     return VisionEncoderDecoderDecoderOnnxConfig(__lowerCAmelCase     ,			__lowerCAmelCase				)












"""simple docstring"""



import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
SCREAMING_SNAKE_CASE    =	object()

# For specifying empty leaf dict `{}`
SCREAMING_SNAKE_CASE    =	object()







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_						)						->		List[Any]:
 A__	=    tuple((re.compile(x + "$"						) for x in qs)						)
 for i in range(len(lowercase_						) - len(lowercase_						) + 1						):
  A__	=    [x.match(lowercase_						) for x, y in zip(lowercase_      ,			ks[i:]						)]
  if matches and all(lowercase_						):
   return True
 return False







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		Optional[int]:
 def replace(lowercase_      ,			lowercase_						):
  for rule, replacement in rules:
   if _match(lowercase_      ,			lowercase_						):
    return replacement
  return val

 return replace







def   _SCREAMING_SNAKE_CASE       (							)						->		List[str]:
 return [
     # embeddings
     (("transformer", "wpe", "embedding"), P("mp"      ,			lowercase_						)),
     (("transformer", "wte", "embedding"), P("mp"      ,			lowercase_						)),
     # atention
     (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(lowercase_      ,			"mp"						)),
     (("attention", "out_proj", "kernel"), P("mp"      ,			lowercase_						)),
     (("attention", "out_proj", "bias"), None),
     # mlp
     (("mlp", "c_fc", "kernel"), P(lowercase_      ,			"mp"						)),
     (("mlp", "c_fc", "bias"), P("mp"						)),
     (("mlp", "c_proj", "kernel"), P("mp"      ,			lowercase_						)),
     (("mlp", "c_proj", "bias"), None),
     # layer norms
     ((r"ln_\d+", "bias"), None),
     ((r"\d+", r"ln_\d+", "scale"), None),
     (("ln_f", "bias"), None),
     (("ln_f", "scale"), None),
 ]







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		int:
 A__	=    _get_partition_rules()
 A__	=    _replacement_rules(lowercase_						)
 A__	=    {k: _unmatched for k in flatten_dict(lowercase_						)}
 A__	=    {k: replace(lowercase_      ,			lowercase_						) for k, v in initd.items()}
 assert _unmatched not in result.values(), "Incomplete partition spec."
 return freeze(unflatten_dict(lowercase_						)						)





