





from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_A					:				Union[str, Any]				=			logging.get_logger(__name__)

_A					:				List[Any]				=			{
    """facebook/xlm-roberta-xl""": """https://huggingface.co/facebook/xlm-roberta-xl/resolve/main/config.json""",
    """facebook/xlm-roberta-xxl""": """https://huggingface.co/facebook/xlm-roberta-xxl/resolve/main/config.json""",
    # See all XLM-RoBERTa-XL models at https://huggingface.co/models?filter=xlm-roberta-xl
}



class 					__snake_case	(  __SCREAMING_SNAKE_CASE    ):


    '''simple docstring'''







    lowerCamelCase__    :		List[str]				=	"""xlm-roberta-xl"""



    def __init__(       self   ,				A_=25_08_80   ,				A_=25_60   ,				A_=36   ,				A_=32   ,				A_=1_02_40   ,				A_="gelu"   ,				A_=0.1   ,				A_=0.1   ,				A_=5_14   ,				A_=1   ,				A_=0.02   ,				A_=1E-05   ,				A_=1   ,				A_=0   ,				A_=2   ,				A_="absolute"   ,				A_=True   ,				A_=None   ,				**A_   ,				):





        '''simple docstring'''

        super().__init__(pad_token_id=A_   ,				bos_token_id=A_   ,				eos_token_id=A_   ,				**A_     )
        SCREAMING_SNAKE_CASE__											=			vocab_size
        SCREAMING_SNAKE_CASE__											=			hidden_size
        SCREAMING_SNAKE_CASE__											=			num_hidden_layers
        SCREAMING_SNAKE_CASE__											=			num_attention_heads
        SCREAMING_SNAKE_CASE__											=			hidden_act
        SCREAMING_SNAKE_CASE__											=			intermediate_size
        SCREAMING_SNAKE_CASE__											=			hidden_dropout_prob
        SCREAMING_SNAKE_CASE__											=			attention_probs_dropout_prob
        SCREAMING_SNAKE_CASE__											=			max_position_embeddings
        SCREAMING_SNAKE_CASE__											=			type_vocab_size
        SCREAMING_SNAKE_CASE__											=			initializer_range
        SCREAMING_SNAKE_CASE__											=			layer_norm_eps
        SCREAMING_SNAKE_CASE__											=			position_embedding_type
        SCREAMING_SNAKE_CASE__											=			use_cache
        SCREAMING_SNAKE_CASE__											=			classifier_dropout





class 					__snake_case	(  __SCREAMING_SNAKE_CASE    ):


    '''simple docstring'''







    @property
    def  lowercase_					(       self     ):





        '''simple docstring'''

        if self.task == "multiple-choice":
            SCREAMING_SNAKE_CASE__											=			{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
        else:
            SCREAMING_SNAKE_CASE__											=			{0: '''batch''', 1: '''sequence'''}
        return OrderedDict(
            [
                ('''input_ids''', dynamic_axis),
                ('''attention_mask''', dynamic_axis),
            ]     )




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

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD






class        lowerCAmelCase_  (     nn.Module	):
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=0.0
       __UpperCAmelCase					=1
       __UpperCAmelCase					=1
       __UpperCAmelCase					=True
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=jnp.floataa






       def   UpperCamelCase						( self   )->      int:
            _A				=							[]
            _A				=							[]

            for i in range(self.num_layers   ):
                 _A				=							self.in_channels if i == 0 else self.out_channels

                 _A				=							FlaxResnetBlockaD(
                     in_channels=_UpperCamelCase			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
                 resnets.append(_UpperCamelCase   )

                 _A				=							FlaxTransformeraDModel(
                     in_channels=self.out_channels			,     n_heads=self.num_attention_heads			,     d_head=self.out_channels // self.num_attention_heads			,     depth=1			,     use_linear_projection=self.use_linear_projection			,     only_cross_attention=self.only_cross_attention			,     use_memory_efficient_attention=self.use_memory_efficient_attention			,     dtype=self.dtype			,     )
                 attentions.append(_UpperCamelCase   )

            _A				=							resnets
            _A				=							attentions

            if self.add_downsample:
                 _A				=							FlaxDownsampleaD(self.out_channels			,     dtype=self.dtype   )






       def __call__( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase=True   )->      Dict:
            _A				=							()

            for resnet, attn in zip(self.resnets			,     self.attentions   ):
                 _A				=							resnet(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )
                 _A				=							attn(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )
                 output_states += (hidden_states,)

            if self.add_downsample:
                 _A				=							self.downsamplers_a(_UpperCamelCase   )
                 output_states += (hidden_states,)

            return hidden_states, output_states






class        lowerCAmelCase_  (     nn.Module	):
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=0.0
       __UpperCAmelCase					=1
       __UpperCAmelCase					=True
       __UpperCAmelCase					=jnp.floataa






       def   UpperCamelCase						( self   )->      Optional[Any]:
            _A				=							[]

            for i in range(self.num_layers   ):
                 _A				=							self.in_channels if i == 0 else self.out_channels

                 _A				=							FlaxResnetBlockaD(
                     in_channels=_UpperCamelCase			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
                 resnets.append(_UpperCamelCase   )
            _A				=							resnets

            if self.add_downsample:
                 _A				=							FlaxDownsampleaD(self.out_channels			,     dtype=self.dtype   )






       def __call__( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase=True   )->      Union[str, Any]:
            _A				=							()

            for resnet in self.resnets:
                 _A				=							resnet(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )
                 output_states += (hidden_states,)

            if self.add_downsample:
                 _A				=							self.downsamplers_a(_UpperCamelCase   )
                 output_states += (hidden_states,)

            return hidden_states, output_states






class        lowerCAmelCase_  (     nn.Module	):
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=0.0
       __UpperCAmelCase					=1
       __UpperCAmelCase					=1
       __UpperCAmelCase					=True
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=jnp.floataa






       def   UpperCamelCase						( self   )->      str:
            _A				=							[]
            _A				=							[]

            for i in range(self.num_layers   ):
                 _A				=							self.in_channels if (i == self.num_layers - 1) else self.out_channels
                 _A				=							self.prev_output_channel if i == 0 else self.out_channels

                 _A				=							FlaxResnetBlockaD(
                     in_channels=resnet_in_channels + res_skip_channels			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
                 resnets.append(_UpperCamelCase   )

                 _A				=							FlaxTransformeraDModel(
                     in_channels=self.out_channels			,     n_heads=self.num_attention_heads			,     d_head=self.out_channels // self.num_attention_heads			,     depth=1			,     use_linear_projection=self.use_linear_projection			,     only_cross_attention=self.only_cross_attention			,     use_memory_efficient_attention=self.use_memory_efficient_attention			,     dtype=self.dtype			,     )
                 attentions.append(_UpperCamelCase   )

            _A				=							resnets
            _A				=							attentions

            if self.add_upsample:
                 _A				=							FlaxUpsampleaD(self.out_channels			,     dtype=self.dtype   )






       def __call__( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase=True   )->      Dict:
            for resnet, attn in zip(self.resnets			,     self.attentions   ):
                 # pop res hidden states
                 _A				=							res_hidden_states_tuple[-1]
                 _A				=							res_hidden_states_tuple[:-1]
                 _A				=							jnp.concatenate((hidden_states, res_hidden_states)			,     axis=-1   )

                 _A				=							resnet(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )
                 _A				=							attn(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )

            if self.add_upsample:
                 _A				=							self.upsamplers_a(_UpperCamelCase   )

            return hidden_states






class        lowerCAmelCase_  (     nn.Module	):
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=0.0
       __UpperCAmelCase					=1
       __UpperCAmelCase					=True
       __UpperCAmelCase					=jnp.floataa






       def   UpperCamelCase						( self   )->      Optional[Any]:
            _A				=							[]

            for i in range(self.num_layers   ):
                 _A				=							self.in_channels if (i == self.num_layers - 1) else self.out_channels
                 _A				=							self.prev_output_channel if i == 0 else self.out_channels

                 _A				=							FlaxResnetBlockaD(
                     in_channels=resnet_in_channels + res_skip_channels			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
                 resnets.append(_UpperCamelCase   )

            _A				=							resnets

            if self.add_upsample:
                 _A				=							FlaxUpsampleaD(self.out_channels			,     dtype=self.dtype   )






       def __call__( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase=True   )->      Dict:
            for resnet in self.resnets:
                 # pop res hidden states
                 _A				=							res_hidden_states_tuple[-1]
                 _A				=							res_hidden_states_tuple[:-1]
                 _A				=							jnp.concatenate((hidden_states, res_hidden_states)			,     axis=-1   )

                 _A				=							resnet(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )

            if self.add_upsample:
                 _A				=							self.upsamplers_a(_UpperCamelCase   )

            return hidden_states






class        lowerCAmelCase_  (     nn.Module	):
       __UpperCAmelCase					=42
       __UpperCAmelCase					=0.0
       __UpperCAmelCase					=1
       __UpperCAmelCase					=1
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=jnp.floataa






       def   UpperCamelCase						( self   )->      Dict:
            # there is always at least one resnet
            _A				=							[
                FlaxResnetBlockaD(
                    in_channels=self.in_channels			,     out_channels=self.in_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
            ]

            _A				=							[]

            for _ in range(self.num_layers   ):
                 _A				=							FlaxTransformeraDModel(
                     in_channels=self.in_channels			,     n_heads=self.num_attention_heads			,     d_head=self.in_channels // self.num_attention_heads			,     depth=1			,     use_linear_projection=self.use_linear_projection			,     use_memory_efficient_attention=self.use_memory_efficient_attention			,     dtype=self.dtype			,     )
                 attentions.append(_UpperCamelCase   )

                 _A				=							FlaxResnetBlockaD(
                     in_channels=self.in_channels			,     out_channels=self.in_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
                 resnets.append(_UpperCamelCase   )

            _A				=							resnets
            _A				=							attentions






       def __call__( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase=True   )->      List[Any]:
            _A				=							self.resnets[0](_UpperCamelCase			,     _UpperCamelCase   )
            for attn, resnet in zip(self.attentions			,     self.resnets[1:]   ):
                 _A				=							attn(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )
                 _A				=							resnet(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )

            return hidden_states





