






"""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   UpperCamelCase_ (       nn.Module  ):
  _A       :					int
  _A       :					int
  _A       :					float				=	0.0
  _A       :					int				=	1
  _A       :					int				=	1
  _A       :					bool				=	True
  _A       :					bool				=	False
  _A       :					bool				=	False
  _A       :					bool				=	False
  _A       :					jnp.dtype				=	jnp.floataa



  def 		UpperCamelCase_				(      self				)			->			Union[str, Any]:



         """simple docstring"""





         UpperCAmelCase	    =			[]
         UpperCAmelCase	    =			[]

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

                UpperCAmelCase	    =			FlaxResnetBlockaD(
                    in_channels=snake_case__ ,		out_channels=self.out_channels ,		dropout_prob=self.dropout ,		dtype=self.dtype ,		)
                resnets.append(snake_case__				)

                UpperCAmelCase	    =			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(snake_case__				)

         UpperCAmelCase	    =			resnets
         UpperCAmelCase	    =			attentions

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



  def __call__(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__=True				)			->			Any:



         """simple docstring"""





         UpperCAmelCase	    =			()

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

         if self.add_downsample:
                UpperCAmelCase	    =			self.downsamplers_a(snake_case__				)
                output_states += (hidden_states,)

         return hidden_states, output_states





class   UpperCamelCase_ (       nn.Module  ):
  _A       :					int
  _A       :					int
  _A       :					float				=	0.0
  _A       :					int				=	1
  _A       :					bool				=	True
  _A       :					jnp.dtype				=	jnp.floataa



  def 		UpperCamelCase_				(      self				)			->			Optional[Any]:



         """simple docstring"""





         UpperCAmelCase	    =			[]

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

                UpperCAmelCase	    =			FlaxResnetBlockaD(
                    in_channels=snake_case__ ,		out_channels=self.out_channels ,		dropout_prob=self.dropout ,		dtype=self.dtype ,		)
                resnets.append(snake_case__				)
         UpperCAmelCase	    =			resnets

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



  def __call__(      self ,		snake_case__ ,		snake_case__ ,		snake_case__=True				)			->			List[Any]:



         """simple docstring"""





         UpperCAmelCase	    =			()

         for resnet in self.resnets:
                UpperCAmelCase	    =			resnet(snake_case__ ,		snake_case__ ,		deterministic=snake_case__				)
                output_states += (hidden_states,)

         if self.add_downsample:
                UpperCAmelCase	    =			self.downsamplers_a(snake_case__				)
                output_states += (hidden_states,)

         return hidden_states, output_states





class   UpperCamelCase_ (       nn.Module  ):
  _A       :					int
  _A       :					int
  _A       :					int
  _A       :					float				=	0.0
  _A       :					int				=	1
  _A       :					int				=	1
  _A       :					bool				=	True
  _A       :					bool				=	False
  _A       :					bool				=	False
  _A       :					bool				=	False
  _A       :					jnp.dtype				=	jnp.floataa



  def 		UpperCamelCase_				(      self				)			->			List[Any]:



         """simple docstring"""





         UpperCAmelCase	    =			[]
         UpperCAmelCase	    =			[]

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

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

                UpperCAmelCase	    =			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(snake_case__				)

         UpperCAmelCase	    =			resnets
         UpperCAmelCase	    =			attentions

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



  def __call__(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__=True				)			->			Optional[int]:



         """simple docstring"""





         for resnet, attn in zip(self.resnets ,		self.attentions				):
                # pop res hidden states
                UpperCAmelCase	    =			res_hidden_states_tuple[-1]
                UpperCAmelCase	    =			res_hidden_states_tuple[:-1]
                UpperCAmelCase	    =			jnp.concatenate((hidden_states, res_hidden_states) ,		axis=-1				)

                UpperCAmelCase	    =			resnet(snake_case__ ,		snake_case__ ,		deterministic=snake_case__				)
                UpperCAmelCase	    =			attn(snake_case__ ,		snake_case__ ,		deterministic=snake_case__				)

         if self.add_upsample:
                UpperCAmelCase	    =			self.upsamplers_a(snake_case__				)

         return hidden_states





class   UpperCamelCase_ (       nn.Module  ):
  _A       :					int
  _A       :					int
  _A       :					int
  _A       :					float				=	0.0
  _A       :					int				=	1
  _A       :					bool				=	True
  _A       :					jnp.dtype				=	jnp.floataa



  def 		UpperCamelCase_				(      self				)			->			Tuple:



         """simple docstring"""





         UpperCAmelCase	    =			[]

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

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

         UpperCAmelCase	    =			resnets

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



  def __call__(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__=True				)			->			List[str]:



         """simple docstring"""





         for resnet in self.resnets:
                # pop res hidden states
                UpperCAmelCase	    =			res_hidden_states_tuple[-1]
                UpperCAmelCase	    =			res_hidden_states_tuple[:-1]
                UpperCAmelCase	    =			jnp.concatenate((hidden_states, res_hidden_states) ,		axis=-1				)

                UpperCAmelCase	    =			resnet(snake_case__ ,		snake_case__ ,		deterministic=snake_case__				)

         if self.add_upsample:
                UpperCAmelCase	    =			self.upsamplers_a(snake_case__				)

         return hidden_states





class   UpperCamelCase_ (       nn.Module  ):
  _A       :					int
  _A       :					float				=	0.0
  _A       :					int				=	1
  _A       :					int				=	1
  _A       :					bool				=	False
  _A       :					bool				=	False
  _A       :					jnp.dtype				=	jnp.floataa



  def 		UpperCamelCase_				(      self				)			->			str:



         """simple docstring"""





         UpperCAmelCase	    =			[
             FlaxResnetBlockaD(
                 in_channels=self.in_channels ,		out_channels=self.in_channels ,		dropout_prob=self.dropout ,		dtype=self.dtype ,		)
         ]

         UpperCAmelCase	    =			[]

         for _ in range(self.num_layers				):
                UpperCAmelCase	    =			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(snake_case__				)

                UpperCAmelCase	    =			FlaxResnetBlockaD(
                    in_channels=self.in_channels ,		out_channels=self.in_channels ,		dropout_prob=self.dropout ,		dtype=self.dtype ,		)
                resnets.append(snake_case__				)

         UpperCAmelCase	    =			resnets
         UpperCAmelCase	    =			attentions



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



         """simple docstring"""





         UpperCAmelCase	    =			self.resnets[0](snake_case__ ,		snake_case__				)
         for attn, resnet in zip(self.attentions ,		self.resnets[1:]				):
                UpperCAmelCase	    =			attn(snake_case__ ,		snake_case__ ,		deterministic=snake_case__				)
                UpperCAmelCase	    =			resnet(snake_case__ ,		snake_case__ ,		deterministic=snake_case__				)

         return hidden_states








"""simple docstring"""







from collections import OrderedDict
from typing import Mapping

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


lowerCAmelCase_	:    int            =      logging.get_logger(__name__)

lowerCAmelCase_	:    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   UpperCamelCase_ (       a_  ):
  _A       :					int				=	'xlm-roberta-xl'



  def __init__(      self ,		snake_case__=25_08_80 ,		snake_case__=25_60 ,		snake_case__=36 ,		snake_case__=32 ,		snake_case__=1_02_40 ,		snake_case__="gelu" ,		snake_case__=0.1 ,		snake_case__=0.1 ,		snake_case__=5_14 ,		snake_case__=1 ,		snake_case__=0.02 ,		snake_case__=1e-05 ,		snake_case__=1 ,		snake_case__=0 ,		snake_case__=2 ,		snake_case__="absolute" ,		snake_case__=True ,		snake_case__=None ,		**snake_case__ ,		)			->			Union[str, Any]:



         """simple docstring"""





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





class   UpperCamelCase_ (       a_  ):



  @property
  def 		UpperCamelCase_				(      self				)			->			Mapping[str, Mapping[int, str]]:



         """simple docstring"""





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