




'''simple docstring'''
import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp
def 							snake_case_    (  a__						:      str       ,a__						:      Optional[int]       ,a__						:      Any       ,a__						:      Optional[int]       ,a__						:      int = 40_96     ):



    """simple docstring"""




    __lowercase							,__lowercase							,__lowercase     							=		key.shape[-3:]
    __lowercase     							=		value.shape[-1]
    __lowercase     							=		min(UpperCamelCase__       ,UpperCamelCase__     )
    __lowercase     							=		query / jnp.sqrt(UpperCamelCase__     )

    @functools.partial(jax.checkpoint       ,prevent_cse=UpperCamelCase__     )
    def summarize_chunk(a__						:      Union[str, Any]       ,a__						:      Union[str, Any]       ,a__						:      Optional[int]     ):
        __lowercase     							=		jnp.einsum("""...qhd,...khd->...qhk"""       ,UpperCamelCase__       ,UpperCamelCase__       ,precision=UpperCamelCase__     )

        __lowercase     							=		jnp.max(UpperCamelCase__       ,axis=-1       ,keepdims=UpperCamelCase__     )
        __lowercase     							=		jax.lax.stop_gradient(UpperCamelCase__     )
        __lowercase     							=		jnp.exp(attn_weights - max_score     )

        __lowercase     							=		jnp.einsum("""...vhf,...qhv->...qhf"""       ,UpperCamelCase__       ,UpperCamelCase__       ,precision=UpperCamelCase__     )
        __lowercase     							=		jnp.einsum("""...qhk->...qh"""       ,UpperCamelCase__     )

        return (exp_values, exp_weights.sum(axis=-1     ), max_score)

    def chunk_scanner(a__						:      Tuple     ):
        # julienne key array
        __lowercase     							=		jax.lax.dynamic_slice(
            operand=UpperCamelCase__       ,start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]       ,slice_sizes=list(key.shape[:-3]     ) + [key_chunk_size, num_heads, k_features]       ,)

        # julienne value array
        __lowercase     							=		jax.lax.dynamic_slice(
            operand=UpperCamelCase__       ,start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]       ,slice_sizes=list(value.shape[:-3]     ) + [key_chunk_size, num_heads, v_features]       ,)

        return summarize_chunk(UpperCamelCase__       ,UpperCamelCase__       ,UpperCamelCase__     )

    __lowercase							,__lowercase							,__lowercase     							=		jax.lax.map(f=UpperCamelCase__       ,xs=jnp.arange(0       ,UpperCamelCase__       ,UpperCamelCase__     )     )

    __lowercase     							=		jnp.max(UpperCamelCase__       ,axis=0       ,keepdims=UpperCamelCase__     )
    __lowercase     							=		jnp.exp(chunk_max - global_max     )

    chunk_values *= jnp.expand_dims(UpperCamelCase__       ,axis=-1     )
    chunk_weights *= max_diffs

    __lowercase     							=		chunk_values.sum(axis=0     )
    __lowercase     							=		jnp.expand_dims(UpperCamelCase__       ,-1     ).sum(axis=0     )

    return all_values / all_weights
def 							snake_case_    (  a__						:      Tuple       ,a__						:      Optional[Any]       ,a__						:      Dict       ,a__						:      Union[str, Any]=jax.lax.Precision.HIGHEST       ,a__						:      int = 10_24       ,a__						:      int = 40_96     ):



    """simple docstring"""




    __lowercase							,__lowercase							,__lowercase     							=		query.shape[-3:]

    def chunk_scanner(a__						:      str       ,a__						:      Tuple     ):
        # julienne query array
        __lowercase     							=		jax.lax.dynamic_slice(
            operand=UpperCamelCase__       ,start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]       ,slice_sizes=list(query.shape[:-3]     ) + [min(UpperCamelCase__       ,UpperCamelCase__     ), num_heads, q_features]       ,)

        return (
            chunk_idx + query_chunk_size,  # unused ignore it
            _query_chunk_attention(
                query=UpperCamelCase__       ,key=UpperCamelCase__       ,value=UpperCamelCase__       ,precision=UpperCamelCase__       ,key_chunk_size=UpperCamelCase__     ),
        )

    __lowercase							,__lowercase     							=		jax.lax.scan(
        f=UpperCamelCase__       ,init=0       ,xs=UpperCamelCase__       ,length=math.ceil(num_q / query_chunk_size     )     )

    return jnp.concatenate(UpperCamelCase__       ,axis=-3     )  # fuse the chunked result back






class 			SCREAMING_SNAKE_CASE(							nn.Module	):
       snake_case_					:	int
       snake_case_					:	int       =   8
       snake_case_					:	int       =   64
       snake_case_					:	float       =   0.0
       snake_case_					:	bool       =   False
       snake_case_					:	jnp.dtype       =   jnp.floataa

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

           """simple docstring"""



           __lowercase     							=		self.dim_head * self.heads
           __lowercase     							=		self.dim_head**-0.5

           # Weights were exported with old names {to_q, to_k, to_v, to_out}
           __lowercase     							=		nn.Dense(__UpperCAmelCase   ,						use_bias=__UpperCAmelCase   ,						dtype=self.dtype   ,						name="""to_q"""	)
           __lowercase     							=		nn.Dense(__UpperCAmelCase   ,						use_bias=__UpperCAmelCase   ,						dtype=self.dtype   ,						name="""to_k"""	)
           __lowercase     							=		nn.Dense(__UpperCAmelCase   ,						use_bias=__UpperCAmelCase   ,						dtype=self.dtype   ,						name="""to_v"""	)

           __lowercase     							=		nn.Dense(self.query_dim   ,						dtype=self.dtype   ,						name="""to_out_0"""	)
           __lowercase     							=		nn.Dropout(rate=self.dropout	)

       def        snake_case__     (       self   ,						lowerCamelCase__	)       ->   Dict:

           """simple docstring"""



           __lowercase							,__lowercase							,__lowercase     							=		tensor.shape
           __lowercase     							=		self.heads
           __lowercase     							=		tensor.reshape(__UpperCAmelCase   ,						__UpperCAmelCase   ,						__UpperCAmelCase   ,						dim // head_size	)
           __lowercase     							=		jnp.transpose(__UpperCAmelCase   ,						(0, 2, 1, 3)	)
           __lowercase     							=		tensor.reshape(batch_size * head_size   ,						__UpperCAmelCase   ,						dim // head_size	)
           return tensor

       def        snake_case__     (       self   ,						lowerCamelCase__	)       ->   List[Any]:

           """simple docstring"""



           __lowercase							,__lowercase							,__lowercase     							=		tensor.shape
           __lowercase     							=		self.heads
           __lowercase     							=		tensor.reshape(batch_size // head_size   ,						__UpperCAmelCase   ,						__UpperCAmelCase   ,						__UpperCAmelCase	)
           __lowercase     							=		jnp.transpose(__UpperCAmelCase   ,						(0, 2, 1, 3)	)
           __lowercase     							=		tensor.reshape(batch_size // head_size   ,						__UpperCAmelCase   ,						dim * head_size	)
           return tensor




       def __call__(       self   ,						lowerCamelCase__   ,						lowerCamelCase__=None   ,						lowerCamelCase__=True	)       ->   Dict:

           """simple docstring"""



           __lowercase     							=		hidden_states if context is None else context

           __lowercase     							=		self.query(__UpperCAmelCase	)
           __lowercase     							=		self.key(__UpperCAmelCase	)
           __lowercase     							=		self.value(__UpperCAmelCase	)

           __lowercase     							=		self.reshape_heads_to_batch_dim(__UpperCAmelCase	)
           __lowercase     							=		self.reshape_heads_to_batch_dim(__UpperCAmelCase	)
           __lowercase     							=		self.reshape_heads_to_batch_dim(__UpperCAmelCase	)

           if self.use_memory_efficient_attention:
               __lowercase     							=		query_states.transpose(1   ,						0   ,						2	)
               __lowercase     							=		key_states.transpose(1   ,						0   ,						2	)
               __lowercase     							=		value_states.transpose(1   ,						0   ,						2	)

               # this if statement create a chunk size for each layer of the unet
               # the chunk size is equal to the query_length dimension of the deepest layer of the unet

               __lowercase     							=		query_states.shape[-3]
               if flatten_latent_dim % 64 == 0:
                   __lowercase     							=		int(flatten_latent_dim / 64	)
               elif flatten_latent_dim % 16 == 0:
                   __lowercase     							=		int(flatten_latent_dim / 16	)
               elif flatten_latent_dim % 4 == 0:
                   __lowercase     							=		int(flatten_latent_dim / 4	)
               else:
                   __lowercase     							=		int(__UpperCAmelCase	)

               __lowercase     							=		jax_memory_efficient_attention(
                   __UpperCAmelCase   ,						__UpperCAmelCase   ,						__UpperCAmelCase   ,						query_chunk_size=__UpperCAmelCase   ,						key_chunk_size=4096 * 4	)

               __lowercase     							=		hidden_states.transpose(1   ,						0   ,						2	)
           else:
               # compute attentions
               __lowercase     							=		jnp.einsum("""b i d, b j d->b i j"""   ,						__UpperCAmelCase   ,						__UpperCAmelCase	)
               __lowercase     							=		attention_scores * self.scale
               __lowercase     							=		nn.softmax(__UpperCAmelCase   ,						axis=2	)

               # attend to values
               __lowercase     							=		jnp.einsum("""b i j, b j d -> b i d"""   ,						__UpperCAmelCase   ,						__UpperCAmelCase	)

           __lowercase     							=		self.reshape_batch_dim_to_heads(__UpperCAmelCase	)
           __lowercase     							=		self.proj_attn(__UpperCAmelCase	)
           return self.dropout_layer(__UpperCAmelCase   ,						deterministic=__UpperCAmelCase	)






class 			SCREAMING_SNAKE_CASE(							nn.Module	):
       snake_case_					:	int
       snake_case_					:	int
       snake_case_					:	int
       snake_case_					:	float       =   0.0
       snake_case_					:	bool       =   False
       snake_case_					:	jnp.dtype       =   jnp.floataa
       snake_case_					:	bool       =   False

       def        snake_case__     (       self	)       ->   str:

           """simple docstring"""



           __lowercase     							=		FlaxAttention(
               self.dim   ,						self.n_heads   ,						self.d_head   ,						self.dropout   ,						self.use_memory_efficient_attention   ,						dtype=self.dtype	)
           # cross attention
           __lowercase     							=		FlaxAttention(
               self.dim   ,						self.n_heads   ,						self.d_head   ,						self.dropout   ,						self.use_memory_efficient_attention   ,						dtype=self.dtype	)
           __lowercase     							=		FlaxFeedForward(dim=self.dim   ,						dropout=self.dropout   ,						dtype=self.dtype	)
           __lowercase     							=		nn.LayerNorm(epsilon=1E-5   ,						dtype=self.dtype	)
           __lowercase     							=		nn.LayerNorm(epsilon=1E-5   ,						dtype=self.dtype	)
           __lowercase     							=		nn.LayerNorm(epsilon=1E-5   ,						dtype=self.dtype	)
           __lowercase     							=		nn.Dropout(rate=self.dropout	)




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

           """simple docstring"""



           __lowercase     							=		hidden_states
           if self.only_cross_attention:
               __lowercase     							=		self.attna(self.norma(__UpperCAmelCase	)   ,						__UpperCAmelCase   ,						deterministic=__UpperCAmelCase	)
           else:
               __lowercase     							=		self.attna(self.norma(__UpperCAmelCase	)   ,						deterministic=__UpperCAmelCase	)
           __lowercase     							=		hidden_states + residual

           # cross attention
           __lowercase     							=		hidden_states
           __lowercase     							=		self.attna(self.norma(__UpperCAmelCase	)   ,						__UpperCAmelCase   ,						deterministic=__UpperCAmelCase	)
           __lowercase     							=		hidden_states + residual

           # feed forward
           __lowercase     							=		hidden_states
           __lowercase     							=		self.ff(self.norma(__UpperCAmelCase	)   ,						deterministic=__UpperCAmelCase	)
           __lowercase     							=		hidden_states + residual

           return self.dropout_layer(__UpperCAmelCase   ,						deterministic=__UpperCAmelCase	)






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

       def        snake_case__     (       self	)       ->   int:

           """simple docstring"""



           __lowercase     							=		nn.GroupNorm(num_groups=32   ,						epsilon=1E-5	)

           __lowercase     							=		self.n_heads * self.d_head
           if self.use_linear_projection:
               __lowercase     							=		nn.Dense(__UpperCAmelCase   ,						dtype=self.dtype	)
           else:
               __lowercase     							=		nn.Conv(
                   __UpperCAmelCase   ,						kernel_size=(1, 1)   ,						strides=(1, 1)   ,						padding="""VALID"""   ,						dtype=self.dtype   ,						)

           __lowercase     							=		[
               FlaxBasicTransformerBlock(
                   __UpperCAmelCase   ,						self.n_heads   ,						self.d_head   ,						dropout=self.dropout   ,						only_cross_attention=self.only_cross_attention   ,						dtype=self.dtype   ,						use_memory_efficient_attention=self.use_memory_efficient_attention   ,						)
               for _ in range(self.depth	)
           ]

           if self.use_linear_projection:
               __lowercase     							=		nn.Dense(__UpperCAmelCase   ,						dtype=self.dtype	)
           else:
               __lowercase     							=		nn.Conv(
                   __UpperCAmelCase   ,						kernel_size=(1, 1)   ,						strides=(1, 1)   ,						padding="""VALID"""   ,						dtype=self.dtype   ,						)

           __lowercase     							=		nn.Dropout(rate=self.dropout	)




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

           """simple docstring"""



           __lowercase							,__lowercase							,__lowercase							,__lowercase     							=		hidden_states.shape
           __lowercase     							=		hidden_states
           __lowercase     							=		self.norm(__UpperCAmelCase	)
           if self.use_linear_projection:
               __lowercase     							=		hidden_states.reshape(__UpperCAmelCase   ,						height * width   ,						__UpperCAmelCase	)
               __lowercase     							=		self.proj_in(__UpperCAmelCase	)
           else:
               __lowercase     							=		self.proj_in(__UpperCAmelCase	)
               __lowercase     							=		hidden_states.reshape(__UpperCAmelCase   ,						height * width   ,						__UpperCAmelCase	)

           for transformer_block in self.transformer_blocks:
               __lowercase     							=		transformer_block(__UpperCAmelCase   ,						__UpperCAmelCase   ,						deterministic=__UpperCAmelCase	)

           if self.use_linear_projection:
               __lowercase     							=		self.proj_out(__UpperCAmelCase	)
               __lowercase     							=		hidden_states.reshape(__UpperCAmelCase   ,						__UpperCAmelCase   ,						__UpperCAmelCase   ,						__UpperCAmelCase	)
           else:
               __lowercase     							=		hidden_states.reshape(__UpperCAmelCase   ,						__UpperCAmelCase   ,						__UpperCAmelCase   ,						__UpperCAmelCase	)
               __lowercase     							=		self.proj_out(__UpperCAmelCase	)

           __lowercase     							=		hidden_states + residual
           return self.dropout_layer(__UpperCAmelCase   ,						deterministic=__UpperCAmelCase	)






class 			SCREAMING_SNAKE_CASE(							nn.Module	):
       snake_case_					:	int
       snake_case_					:	float       =   0.0
       snake_case_					:	jnp.dtype       =   jnp.floataa

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

           """simple docstring"""



           __lowercase     							=		FlaxGEGLU(self.dim   ,						self.dropout   ,						self.dtype	)
           __lowercase     							=		nn.Dense(self.dim   ,						dtype=self.dtype	)




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

           """simple docstring"""



           __lowercase     							=		self.net_a(__UpperCAmelCase   ,						deterministic=__UpperCAmelCase	)
           __lowercase     							=		self.net_a(__UpperCAmelCase	)
           return hidden_states






class 			SCREAMING_SNAKE_CASE(							nn.Module	):
       snake_case_					:	int
       snake_case_					:	float       =   0.0
       snake_case_					:	jnp.dtype       =   jnp.floataa

       def        snake_case__     (       self	)       ->   Optional[int]:

           """simple docstring"""



           __lowercase     							=		self.dim * 4
           __lowercase     							=		nn.Dense(inner_dim * 2   ,						dtype=self.dtype	)
           __lowercase     							=		nn.Dropout(rate=self.dropout	)




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

           """simple docstring"""



           __lowercase     							=		self.proj(__UpperCAmelCase	)
           __lowercase							,__lowercase     							=		jnp.split(__UpperCAmelCase   ,						2   ,						axis=2	)
           return self.dropout_layer(hidden_linear * nn.gelu(__UpperCAmelCase	)   ,						deterministic=__UpperCAmelCase	)





"""simple docstring"""




from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


__lowerCamelCase									=	logging.get_logger(__name__)

__lowerCamelCase									=	{
    "facebook/blenderbot-3B": "https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json",
    # See all Blenderbot models at https://huggingface.co/models?filter=blenderbot
}
class      UpperCamelCase__(      __A			):
						lowerCAmelCase__	:    List[str]    				=		'blenderbot'
						lowerCAmelCase__	:    Union[str, Any]    				=		['past_key_values']
						lowerCAmelCase__	:    Optional[Any]    				=		{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}





						def __init__(							self		,__UpperCAmelCase=80_08		,__UpperCAmelCase=1_28		,__UpperCAmelCase=2		,__UpperCAmelCase=1_02_40		,__UpperCAmelCase=32		,__UpperCAmelCase=24		,__UpperCAmelCase=1_02_40		,__UpperCAmelCase=32		,__UpperCAmelCase=0.0		,__UpperCAmelCase=0.0		,__UpperCAmelCase=True		,__UpperCAmelCase=True		,__UpperCAmelCase="gelu"		,__UpperCAmelCase=25_60		,__UpperCAmelCase=0.1		,__UpperCAmelCase=0.0		,__UpperCAmelCase=0.0		,__UpperCAmelCase=0.0_2		,__UpperCAmelCase=1		,__UpperCAmelCase=False		,__UpperCAmelCase=0		,__UpperCAmelCase=1		,__UpperCAmelCase=2		,__UpperCAmelCase=3		,__UpperCAmelCase=2		,**__UpperCAmelCase		,)       ->					Any:
													A__						  =				vocab_size
													A__						  =				max_position_embeddings
													A__						  =				d_model
													A__						  =				encoder_ffn_dim
													A__						  =				encoder_layers
													A__						  =				encoder_attention_heads
													A__						  =				decoder_ffn_dim
													A__						  =				decoder_layers
													A__						  =				decoder_attention_heads
													A__						  =				dropout
													A__						  =				attention_dropout
													A__						  =				activation_dropout
													A__						  =				activation_function
													A__						  =				init_std
													A__						  =				encoder_layerdrop
													A__						  =				decoder_layerdrop
													A__						  =				use_cache
													A__						  =				encoder_layers
													A__						  =				scale_embedding  # scale factor will be sqrt(d_model) if True

													super().__init__(
													    pad_token_id=__UpperCAmelCase		,bos_token_id=__UpperCAmelCase		,eos_token_id=__UpperCAmelCase		,is_encoder_decoder=__UpperCAmelCase		,decoder_start_token_id=__UpperCAmelCase		,encoder_no_repeat_ngram_size=__UpperCAmelCase		,forced_eos_token_id=__UpperCAmelCase		,**__UpperCAmelCase		,)






class      UpperCamelCase__(      __A			):





						@property
						def 					snake_case__						(							self  )       ->					Mapping[str, Mapping[int, str]]:
													if self.task in ["default", "seq2seq-lm"]:
																				A__						  =				OrderedDict(
																				    [
																				        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
																				        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
																				    ]  )
																				if self.use_past:
																											A__						  =				{0: 'batch'}
																											A__						  =				{0: 'batch', 1: 'past_decoder_sequence + sequence'}
																				else:
																											A__						  =				{0: 'batch', 1: 'decoder_sequence'}
																											A__						  =				{0: 'batch', 1: 'decoder_sequence'}
																				if self.use_past:
																											self.fill_with_past_key_values_(__UpperCAmelCase		,direction='inputs'  )
													elif self.task == "causal-lm":
																				A__						  =				OrderedDict(
																				    [
																				        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
																				        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
																				    ]  )
																				if self.use_past:
																											A__  ,	A__						  =				self.num_layers
																											for i in range(__UpperCAmelCase  ):
																																		A__						  =				{0: 'batch', 2: 'past_sequence + sequence'}
																																		A__						  =				{0: 'batch', 2: 'past_sequence + sequence'}
													else:
																				A__						  =				OrderedDict(
																				    [
																				        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
																				        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
																				        ('decoder_input_ids', {0: 'batch', 1: 'decoder_sequence'}),
																				        ('decoder_attention_mask', {0: 'batch', 1: 'decoder_sequence'}),
																				    ]  )

													return common_inputs





						@property
						# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
						def 					snake_case__						(							self  )       ->					Mapping[str, Mapping[int, str]]:
													if self.task in ["default", "seq2seq-lm"]:
																				A__						  =				super().outputs
													else:
																				A__						  =				super(__UpperCAmelCase		,self  ).outputs
																				if self.use_past:
																											A__  ,	A__						  =				self.num_layers
																											for i in range(__UpperCAmelCase  ):
																																		A__						  =				{0: 'batch', 2: 'past_sequence + sequence'}
																																		A__						  =				{0: 'batch', 2: 'past_sequence + sequence'}
													return common_outputs





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase = -1		,__UpperCAmelCase = -1		,__UpperCAmelCase = False		,__UpperCAmelCase = None		,)       ->					Mapping[str, Any]:
													A__						  =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
													    __UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )
													# Generate decoder inputs
													A__						  =				seq_length if not self.use_past else 1
													A__						  =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
													    __UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )
													A__						  =				{f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
													A__						  =				dict(**__UpperCAmelCase		,**__UpperCAmelCase  )

													if self.use_past:
																				if not is_torch_available():
																											raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'  )
																				else:
																											import torch
																				A__  ,	A__						  =				common_inputs['input_ids'].shape
																				A__						  =				common_inputs['decoder_input_ids'].shape[1]
																				A__  ,	A__						  =				self.num_attention_heads
																				A__						  =				(
																				    batch,
																				    num_encoder_attention_heads,
																				    encoder_seq_length,
																				    self._config.hidden_size // num_encoder_attention_heads,
																				)
																				A__						  =				decoder_seq_length
																				A__						  =				(
																				    batch,
																				    num_decoder_attention_heads,
																				    decoder_past_length,
																				    self._config.hidden_size // num_decoder_attention_heads,
																				)
																				A__						  =				torch.cat(
																				    [common_inputs['decoder_attention_mask'], torch.ones(__UpperCAmelCase		,__UpperCAmelCase  )]		,dim=1  )
																				A__						  =				[]
																				A__  ,	A__						  =				self.num_layers

																				for _ in range(__UpperCAmelCase  ):
																											common_inputs["past_key_values"].append(
																											    (
																											        torch.zeros(__UpperCAmelCase  ),
																											        torch.zeros(__UpperCAmelCase  ),
																											        torch.zeros(__UpperCAmelCase  ),
																											        torch.zeros(__UpperCAmelCase  ),
																											    )  )
													return common_inputs





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase = -1		,__UpperCAmelCase = -1		,__UpperCAmelCase = False		,__UpperCAmelCase = None		,)       ->					Mapping[str, Any]:
													A__						  =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
													    __UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )

													if self.use_past:
																				if not is_torch_available():
																											raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'  )
																				else:
																											import torch
																				A__  ,	A__						  =				common_inputs['input_ids'].shape
																				A__						  =				seqlen
																				A__  ,	A__						  =				self.num_layers
																				A__  ,	A__						  =				self.num_attention_heads
																				A__						  =				(
																				    batch,
																				    num_encoder_attention_heads,
																				    past_key_values_length,
																				    self._config.hidden_size // num_encoder_attention_heads,
																				)
																				A__						  =				common_inputs['attention_mask'].dtype
																				A__						  =				torch.cat(
																				    [common_inputs['attention_mask'], torch.ones(__UpperCAmelCase		,__UpperCAmelCase		,dtype=__UpperCAmelCase  )]		,dim=1  )
																				A__						  =				[
																				    (torch.zeros(__UpperCAmelCase  ), torch.zeros(__UpperCAmelCase  )) for _ in range(__UpperCAmelCase  )
																				]
													return common_inputs





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase = -1		,__UpperCAmelCase = -1		,__UpperCAmelCase = False		,__UpperCAmelCase = None		,)       ->					Mapping[str, Any]:
													# Copied from OnnxConfig.generate_dummy_inputs
													# Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
													# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
													A__						  =				compute_effective_axis_dimension(
													    __UpperCAmelCase		,fixed_dimension=OnnxConfig.default_fixed_batch		,num_token_to_add=0  )

													# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
													A__						  =				tokenizer.num_special_tokens_to_add(__UpperCAmelCase  )
													A__						  =				compute_effective_axis_dimension(
													    __UpperCAmelCase		,fixed_dimension=OnnxConfig.default_fixed_sequence		,num_token_to_add=__UpperCAmelCase  )

													# Generate dummy inputs according to compute batch and sequence
													A__						  =				[' '.join([tokenizer.unk_token]  ) * seq_length] * batch_size
													A__						  =				dict(tokenizer(__UpperCAmelCase		,return_tensors=__UpperCAmelCase  )  )
													return common_inputs





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase = -1		,__UpperCAmelCase = -1		,__UpperCAmelCase = False		,__UpperCAmelCase = None		,)       ->					Mapping[str, Any]:
													if self.task in ["default", "seq2seq-lm"]:
																				A__						  =				self._generate_dummy_inputs_for_default_and_seqaseq_lm(
																				    __UpperCAmelCase		,batch_size=__UpperCAmelCase		,seq_length=__UpperCAmelCase		,is_pair=__UpperCAmelCase		,framework=__UpperCAmelCase  )

													elif self.task == "causal-lm":
																				A__						  =				self._generate_dummy_inputs_for_causal_lm(
																				    __UpperCAmelCase		,batch_size=__UpperCAmelCase		,seq_length=__UpperCAmelCase		,is_pair=__UpperCAmelCase		,framework=__UpperCAmelCase  )
													else:
																				A__						  =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
																				    __UpperCAmelCase		,batch_size=__UpperCAmelCase		,seq_length=__UpperCAmelCase		,is_pair=__UpperCAmelCase		,framework=__UpperCAmelCase  )

													return common_inputs





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					Tuple:
													if self.task in ["default", "seq2seq-lm"]:
																				A__						  =				super()._flatten_past_key_values_(__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )
													else:
																				A__						  =				super(__UpperCAmelCase		,self  )._flatten_past_key_values_(
																				    __UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					int:
													if direction not in ["inputs", "outputs"]:
																				raise ValueError(f'''direction must either be "inputs" or "outputs", but {direction} was given'''  )

													A__						  =				'past_key_values' if direction == 'inputs' else 'present'
													A__  ,	A__						  =				self.num_layers

													A__						  =				'past_encoder_sequence'
													A__						  =				'past_decoder_sequence' if direction == 'inputs' else 'past_decoder_sequence + sequence'

													for i in range(__UpperCAmelCase  ):
																				A__						  =				{0: 'batch', 2: decoder_sequence}
																				A__						  =				{0: 'batch', 2: decoder_sequence}
																				A__						  =				{0: 'batch', 2: encoder_sequence}
																				A__						  =				{0: 'batch', 2: encoder_sequence}





