
"""simple docstring"""






import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__SCREAMING_SNAKE_CASE							=     logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE							=     {
    'google/pix2struct-textcaps-base': (
        'https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json'
    ),
}





class 				a__					(						A__       ):
     UpperCAmelCase__         =      '''pix2struct_text_model'''
     UpperCAmelCase__         =      ['''past_key_values''']
     UpperCAmelCase__         =      {
         '''hidden_size''': '''hidden_size''',
         '''num_attention_heads''': '''num_heads''',
         '''num_hidden_layers''': '''num_layers''',
     }



     def __init__(							self :str      ,  _lowerCamelCase :Dict=50_244      ,  _lowerCamelCase :Any=768      ,  _lowerCamelCase :Union[str, Any]=64      ,  _lowerCamelCase :Optional[Any]=2_048      ,  _lowerCamelCase :str=12      ,  _lowerCamelCase :List[Any]=12      ,  _lowerCamelCase :Tuple=32      ,  _lowerCamelCase :List[Any]=128      ,  _lowerCamelCase :Any=0.1      ,  _lowerCamelCase :Dict=1E-6      ,  _lowerCamelCase :int=1.0      ,  _lowerCamelCase :str="gelu_new"      ,  _lowerCamelCase :Tuple=0      ,  _lowerCamelCase :Dict=False      ,  _lowerCamelCase :Dict=0      ,  _lowerCamelCase :Any=1      ,  _lowerCamelCase :Optional[Any]=False      ,  _lowerCamelCase :Tuple=True      ,  **_lowerCamelCase :Any      ,  ):







      '''simple docstring'''



      UpperCamelCase_     :					Tuple   	=vocab_size
      UpperCamelCase_     :					Tuple   	=hidden_size
      UpperCamelCase_     :					List[str]   	=d_kv
      UpperCamelCase_     :					Union[str, Any]   	=d_ff
      UpperCamelCase_     :					List[Any]   	=num_layers
      UpperCamelCase_     :					Optional[Any]   	=num_heads
      UpperCamelCase_     :					int   	=relative_attention_num_buckets
      UpperCamelCase_     :					List[Any]   	=relative_attention_max_distance
      UpperCamelCase_     :					List[Any]   	=dropout_rate
      UpperCamelCase_     :					Any   	=layer_norm_epsilon
      UpperCamelCase_     :					Tuple   	=initializer_factor
      UpperCamelCase_     :					Union[str, Any]   	=use_cache

      UpperCamelCase_     :					List[str]   	=eos_token_id
      UpperCamelCase_     :					List[str]   	=decoder_start_token_id

      # for backwards compatibility
      UpperCamelCase_     :					Optional[Any]   	=dense_act_fn

      super().__init__(
          pad_token_id=_lowerCamelCase      ,  eos_token_id=_lowerCamelCase      ,  decoder_start_token_id=_lowerCamelCase      ,  tie_word_embeddings=_lowerCamelCase      ,  is_decoder=_lowerCamelCase      ,  **_lowerCamelCase      ,  )




     @classmethod
     def  lowerCamelCase_						(							cls :List[str]      ,  _lowerCamelCase :Union[str, os.PathLike]      ,  **_lowerCamelCase :Optional[int]   ):







      '''simple docstring'''



      cls._set_token_in_kwargs(_lowerCamelCase   )

      UpperCamelCase_	,     UpperCamelCase_     :					Dict   	=cls.get_config_dict(_lowerCamelCase      ,  **_lowerCamelCase   )

      # get the text config dict if we are loading from Pix2StructConfig
      if config_dict.get('model_type'   ) == "pix2struct":
       UpperCamelCase_     :					Optional[int]   	=config_dict['text_config']

      if "model_type" in config_dict and hasattr(cls      ,  'model_type'   ) and config_dict["model_type"] != cls.model_type:
       logger.warning(
           f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
           f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''   )

      return cls.from_dict(_lowerCamelCase      ,  **_lowerCamelCase   )





class 				a__					(						A__       ):
     UpperCAmelCase__         =      '''pix2struct_vision_model'''



     def __init__(							self :str      ,  _lowerCamelCase :str=768      ,  _lowerCamelCase :Dict=768      ,  _lowerCamelCase :Optional[int]=2_048      ,  _lowerCamelCase :List[str]=64      ,  _lowerCamelCase :Optional[Any]=12      ,  _lowerCamelCase :Dict=12      ,  _lowerCamelCase :Optional[Any]="gelu_new"      ,  _lowerCamelCase :Dict=1E-6      ,  _lowerCamelCase :List[Any]=0.0      ,  _lowerCamelCase :List[str]=0.0      ,  _lowerCamelCase :Dict=1E-1_0      ,  _lowerCamelCase :str=1.0      ,  _lowerCamelCase :Any=4_096      ,  _lowerCamelCase :Any=32      ,  _lowerCamelCase :Optional[Any]=128      ,  **_lowerCamelCase :List[str]      ,  ):







      '''simple docstring'''



      super().__init__(**_lowerCamelCase   )

      UpperCamelCase_     :					Any   	=hidden_size
      UpperCamelCase_     :					Tuple   	=patch_embed_hidden_size
      UpperCamelCase_     :					Optional[Any]   	=d_ff
      UpperCamelCase_     :					Union[str, Any]   	=dropout_rate
      UpperCamelCase_     :					str   	=num_hidden_layers
      UpperCamelCase_     :					str   	=num_attention_heads
      UpperCamelCase_     :					List[Any]   	=initializer_range
      UpperCamelCase_     :					Optional[Any]   	=initializer_factor
      UpperCamelCase_     :					List[str]   	=attention_dropout
      UpperCamelCase_     :					Optional[int]   	=layer_norm_eps
      UpperCamelCase_     :					Optional[Any]   	=dense_act_fn
      UpperCamelCase_     :					Dict   	=seq_len
      UpperCamelCase_     :					List[Any]   	=relative_attention_num_buckets
      UpperCamelCase_     :					List[str]   	=relative_attention_max_distance
      UpperCamelCase_     :					List[Any]   	=d_kv




     @classmethod
     def  lowerCamelCase_						(							cls :Any      ,  _lowerCamelCase :Union[str, os.PathLike]      ,  **_lowerCamelCase :Any   ):







      '''simple docstring'''



      cls._set_token_in_kwargs(_lowerCamelCase   )

      UpperCamelCase_	,     UpperCamelCase_     :					Union[str, Any]   	=cls.get_config_dict(_lowerCamelCase      ,  **_lowerCamelCase   )

      # get the vision config dict if we are loading from Pix2StructConfig
      if config_dict.get('model_type'   ) == "pix2struct":
       UpperCamelCase_     :					int   	=config_dict['vision_config']

      if "model_type" in config_dict and hasattr(cls      ,  'model_type'   ) and config_dict["model_type"] != cls.model_type:
       logger.warning(
           f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
           f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''   )

      return cls.from_dict(_lowerCamelCase      ,  **_lowerCamelCase   )





class 				a__					(						A__       ):
     UpperCAmelCase__         =      '''pix2struct'''
     UpperCAmelCase__         =      True



     def __init__(							self :Dict      ,  _lowerCamelCase :Optional[int]=None      ,  _lowerCamelCase :Optional[Any]=None      ,  _lowerCamelCase :int=1.0      ,  _lowerCamelCase :str=0.02      ,  _lowerCamelCase :Any=False      ,  _lowerCamelCase :Optional[int]=False      ,  _lowerCamelCase :int=True      ,  **_lowerCamelCase :Optional[Any]      ,  ):







      '''simple docstring'''



      super().__init__(tie_word_embeddings=_lowerCamelCase      ,  is_encoder_decoder=_lowerCamelCase      ,  **_lowerCamelCase   )

      if text_config is None:
       UpperCamelCase_     :					List[Any]   	={}
       logger.info('text_config is None. Initializing the Pix2StructTextConfig with default values.'   )

      if vision_config is None:
       UpperCamelCase_     :					str   	={}
       logger.info('vision_config is None. Initializing the Pix2StructVisionConfig with default values.'   )

      UpperCamelCase_     :					Any   	=PixaStructTextConfig(**_lowerCamelCase   )
      UpperCamelCase_     :					Any   	=PixaStructVisionConfig(**_lowerCamelCase   )

      UpperCamelCase_     :					int   	=self.text_config.decoder_start_token_id
      UpperCamelCase_     :					Dict   	=self.text_config.pad_token_id
      UpperCamelCase_     :					Any   	=self.text_config.eos_token_id

      UpperCamelCase_     :					int   	=initializer_factor
      UpperCamelCase_     :					Any   	=initializer_range

      UpperCamelCase_     :					Optional[int]   	=self.initializer_range
      UpperCamelCase_     :					Tuple   	=self.initializer_range

      UpperCamelCase_     :					Union[str, Any]   	=is_vqa



     @classmethod
     def  lowerCamelCase_						(							cls :Union[str, Any]      ,  _lowerCamelCase :PixaStructTextConfig      ,  _lowerCamelCase :PixaStructVisionConfig      ,  **_lowerCamelCase :Union[str, Any]   ):







      '''simple docstring'''



      return cls(text_config=text_config.to_dict()      ,  vision_config=vision_config.to_dict()      ,  **_lowerCamelCase   )




     def  lowerCamelCase_						(							self :int   ):







      '''simple docstring'''



      UpperCamelCase_     :					Dict   	=copy.deepcopy(self.__dict__   )
      UpperCamelCase_     :					Union[str, Any]   	=self.text_config.to_dict()
      UpperCamelCase_     :					str   	=self.vision_config.to_dict()
      UpperCamelCase_     :					Tuple   	=self.__class__.model_type
      return output


"""simple docstring"""






import math

import flax.linen as nn
import jax.numpy as jnp




def       A_  (				__lowercase						,	__lowercase						,	__lowercase = 1						,	__lowercase = 1						,	__lowercase = 1.0e4						,	__lowercase = False						,	__lowercase = 1.0						,	):
 assert timesteps.ndim == 1, "Timesteps should be a 1d-array"
 assert embedding_dim % 2 == 0, F'''Embedding dimension {embedding_dim} should be even'''
 UpperCamelCase_     :					Optional[int]   	=float(embedding_dim // 2	)
 UpperCamelCase_     :					Optional[Any]   	=math.log(max_timescale / min_timescale	) / (num_timescales - freq_shift)
 UpperCamelCase_     :					List[Any]   	=min_timescale * jnp.exp(jnp.arange(__lowercase						,	dtype=jnp.floataa	) * -log_timescale_increment	)
 UpperCamelCase_     :					int   	=jnp.expand_dims(__lowercase						,	1	) * jnp.expand_dims(__lowercase						,	0	)

 # scale embeddings
 UpperCamelCase_     :					List[str]   	=scale * emb

 if flip_sin_to_cos:
  UpperCamelCase_     :					Tuple   	=jnp.concatenate([jnp.cos(__lowercase	), jnp.sin(__lowercase	)]						,	axis=1	)
 else:
  UpperCamelCase_     :					Tuple   	=jnp.concatenate([jnp.sin(__lowercase	), jnp.cos(__lowercase	)]						,	axis=1	)
 UpperCamelCase_     :					List[Any]   	=jnp.reshape(__lowercase						,	[jnp.shape(__lowercase	)[0], embedding_dim]	)
 return signal





class 				a__					(						nn.Module       ):
     UpperCAmelCase__         =      32
     UpperCAmelCase__         =      jnp.floataa



     @nn.compact
     def __call__(							self :Optional[Any]      ,  _lowerCamelCase :List[str]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Dict   	=nn.Dense(self.time_embed_dim      ,  dtype=self.dtype      ,  name='linear_1'   )(_lowerCamelCase   )
      UpperCamelCase_     :					Any   	=nn.silu(_lowerCamelCase   )
      UpperCamelCase_     :					Tuple   	=nn.Dense(self.time_embed_dim      ,  dtype=self.dtype      ,  name='linear_2'   )(_lowerCamelCase   )
      return temb





class 				a__					(						nn.Module       ):
     UpperCAmelCase__         =      32
     UpperCAmelCase__         =      False
     UpperCAmelCase__         =      1



     @nn.compact
     def __call__(							self :Union[str, Any]      ,  _lowerCamelCase :Dict   ):







      '''simple docstring'''



      return get_sinusoidal_embeddings(
          _lowerCamelCase      ,  embedding_dim=self.dim      ,  flip_sin_to_cos=self.flip_sin_to_cos      ,  freq_shift=self.freq_shift   )
