






import os
import sys

__a:     Union[str, Any]   					=					os.path.join(os.path.dirname(__file__), '''src''')
sys.path.append(SRC_DIR)


from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForMaskedLM,
    AutoModelForQuestionAnswering,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    add_start_docstrings,
)


__a:     Union[str, Any]   					=					[
    '''torch''',
    '''numpy''',
    '''tokenizers''',
    '''filelock''',
    '''requests''',
    '''tqdm''',
    '''regex''',
    '''sentencepiece''',
    '''sacremoses''',
    '''importlib_metadata''',
    '''huggingface_hub''',
]




@add_start_docstrings(AutoConfig.__doc__ )
def       _SCREAMING_SNAKE_CASE (							*__snake_case   ,      **__snake_case )				->						Union[str, Any]:
       return AutoConfig.from_pretrained(*__snake_case   ,      **__snake_case )




@add_start_docstrings(AutoTokenizer.__doc__ )
def       _SCREAMING_SNAKE_CASE (							*__snake_case   ,      **__snake_case )				->						Any:
       return AutoTokenizer.from_pretrained(*__snake_case   ,      **__snake_case )




@add_start_docstrings(AutoModel.__doc__ )
def       _SCREAMING_SNAKE_CASE (							*__snake_case   ,      **__snake_case )				->						Tuple:
       return AutoModel.from_pretrained(*__snake_case   ,      **__snake_case )




@add_start_docstrings(AutoModelForCausalLM.__doc__ )
def       _SCREAMING_SNAKE_CASE (							*__snake_case   ,      **__snake_case )				->						Tuple:
       return AutoModelForCausalLM.from_pretrained(*__snake_case   ,      **__snake_case )




@add_start_docstrings(AutoModelForMaskedLM.__doc__ )
def       _SCREAMING_SNAKE_CASE (							*__snake_case   ,      **__snake_case )				->						Optional[Any]:
       return AutoModelForMaskedLM.from_pretrained(*__snake_case   ,      **__snake_case )




@add_start_docstrings(AutoModelForSequenceClassification.__doc__ )
def       _SCREAMING_SNAKE_CASE (							*__snake_case   ,      **__snake_case )				->						List[str]:
       return AutoModelForSequenceClassification.from_pretrained(*__snake_case   ,      **__snake_case )





@add_start_docstrings(AutoModelForQuestionAnswering.__doc__ )
def       _SCREAMING_SNAKE_CASE (							*__snake_case   ,      **__snake_case )				->						List[Any]:
       return AutoModelForQuestionAnswering.from_pretrained(*__snake_case   ,      **__snake_case )







'''simple docstring'''
import flax.linen as nn
import jax
import jax.numpy as jnp


class 							_UpperCAmelCase    (		nn.Module      ):
    __lowerCamelCase: int
    __lowerCamelCase: jnp.dtype									=     jnp.floataa




    def 					lowerCAmelCase__					(      self	:							Optional[Any]      ):






       '''simple docstring'''






       lowercase_						:     List[Any]							     =  nn.Conv(
           self.out_channels	,  kernel_size=(3, 3)	,  strides=(1, 1)	,  padding=((1, 1), (1, 1))	,  dtype=self.dtype	,  )





    def __call__(      self	:							List[str]	,  a	:							Optional[int]      ):






       '''simple docstring'''






       lowercase_     ,		lowercase_     ,		lowercase_     ,		lowercase_						:     Optional[Any]							     =  hidden_states.shape
       lowercase_						:     Tuple							     =  jax.image.resize(
           a	,  shape=(batch, height * 2, width * 2, channels)	,  method="nearest"	,  )
       lowercase_						:     List[Any]							     =  self.conv(a      )
       return hidden_states


class 							_UpperCAmelCase    (		nn.Module      ):
    __lowerCamelCase: int
    __lowerCamelCase: jnp.dtype									=     jnp.floataa




    def 					lowerCAmelCase__					(      self	:							Optional[Any]      ):






       '''simple docstring'''






       lowercase_						:     Any							     =  nn.Conv(
           self.out_channels	,  kernel_size=(3, 3)	,  strides=(2, 2)	,  padding=((1, 1), (1, 1))	,  dtype=self.dtype	,  )





    def __call__(      self	:							Any	,  a	:							int      ):






       '''simple docstring'''






       lowercase_						:     Any							     =  self.conv(a      )
       return hidden_states


class 							_UpperCAmelCase    (		nn.Module      ):
    __lowerCamelCase: int
    __lowerCamelCase: int									=     None
    __lowerCamelCase: float									=     0.0
    __lowerCamelCase: bool									=     None
    __lowerCamelCase: jnp.dtype									=     jnp.floataa




    def 					lowerCAmelCase__					(      self	:							Optional[int]      ):






       '''simple docstring'''






       lowercase_						:     Union[str, Any]							     =  self.in_channels if self.out_channels is None else self.out_channels

       lowercase_						:     Optional[int]							     =  nn.GroupNorm(num_groups=3_2	,  epsilon=1e-5      )
       lowercase_						:     Tuple							     =  nn.Conv(
           a	,  kernel_size=(3, 3)	,  strides=(1, 1)	,  padding=((1, 1), (1, 1))	,  dtype=self.dtype	,  )

       lowercase_						:     List[str]							     =  nn.Dense(a	,  dtype=self.dtype      )

       lowercase_						:     Optional[int]							     =  nn.GroupNorm(num_groups=3_2	,  epsilon=1e-5      )
       lowercase_						:     Any							     =  nn.Dropout(self.dropout_prob      )
       lowercase_						:     Dict							     =  nn.Conv(
           a	,  kernel_size=(3, 3)	,  strides=(1, 1)	,  padding=((1, 1), (1, 1))	,  dtype=self.dtype	,  )

       lowercase_						:     Tuple							     =  self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

       lowercase_						:     Optional[Any]							     =  None
       if use_nin_shortcut:
          lowercase_						:     Union[str, Any]							     =  nn.Conv(
              a	,  kernel_size=(1, 1)	,  strides=(1, 1)	,  padding="VALID"	,  dtype=self.dtype	,  )





    def __call__(      self	:							List[str]	,  a	:							str	,  a	:							Dict	,  a	:							List[str]=True      ):






       '''simple docstring'''






       lowercase_						:     Dict							     =  hidden_states
       lowercase_						:     int							     =  self.norma(a      )
       lowercase_						:     List[Any]							     =  nn.swish(a      )
       lowercase_						:     Dict							     =  self.conva(a      )

       lowercase_						:     Optional[int]							     =  self.time_emb_proj(nn.swish(a      )      )
       lowercase_						:     Tuple							     =  jnp.expand_dims(jnp.expand_dims(a	,  1      )	,  1      )
       lowercase_						:     List[str]							     =  hidden_states + temb

       lowercase_						:     Optional[Any]							     =  self.norma(a      )
       lowercase_						:     Any							     =  nn.swish(a      )
       lowercase_						:     List[str]							     =  self.dropout(a	,  a      )
       lowercase_						:     int							     =  self.conva(a      )

       if self.conv_shortcut is not None:
          lowercase_						:     List[str]							     =  self.conv_shortcut(a      )

       return hidden_states + residual




