




import argparse
import json
import os.path
from collections import OrderedDict

import numpy as np
import requests
import torch
from flax.training.checkpoints import restore_checkpoint
from huggingface_hub import hf_hub_download

from transformers import VivitConfig, VivitForVideoClassification, VivitImageProcessor
from transformers.image_utils import PILImageResampling




def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						Dict      )  ->							Union[str, Any]:
    UpperCAmelCase_							 =     '''https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'''

    with open(__UpperCamelCase       ,     '''wb'''      ) as f:
        with requests.get(__UpperCamelCase       ,     stream=__UpperCamelCase      ) as req:
            for chunk in req.iter_content(chunk_size=2048      ):
                f.write(__UpperCamelCase      )




def 					SCREAMING_SNAKE_CASE	(       )  ->							VivitConfig:
    UpperCAmelCase_							 =     VivitConfig()

    UpperCAmelCase_							 =     400
    UpperCAmelCase_							 =     '''huggingface/label-files'''
    UpperCAmelCase_							 =     '''kinetics400-id2label.json'''

    UpperCAmelCase_							 =     json.load(open(hf_hub_download(__UpperCamelCase       ,     __UpperCamelCase       ,     repo_type='''dataset'''      )       ,     '''r'''      )      )
    UpperCAmelCase_							 =     {int(__UpperCamelCase      ): v for k, v in idalabel.items()}
    UpperCAmelCase_							 =     idalabel
    UpperCAmelCase_							 =     {v: k for k, v in idalabel.items()}
    return config




def 					SCREAMING_SNAKE_CASE	(       )  ->							Union[str, Any]:
    UpperCAmelCase_							 =     hf_hub_download(
        repo_id='''hf-internal-testing/spaghetti-video'''       ,     filename='''eating_spaghetti_32_frames.npy'''       ,     repo_type='''dataset'''      )
    UpperCAmelCase_							 =     np.load(__UpperCamelCase      )
    return list(__UpperCamelCase      )




def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						np.ndarray      )  ->							Optional[Any]:
    if np.ndim(__UpperCamelCase      ) == 2:
        return transform_attention_bias(__UpperCamelCase      )

    elif np.ndim(__UpperCamelCase      ) == 3:
        return transform_attention_kernel(__UpperCamelCase      )

    else:
        raise Exception(f'Invalid number of dimesions: {np.ndim(__UpperCamelCase      )}'      )




def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						np.ndarray      )  ->							Tuple:
    return current.flatten()




def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						np.ndarray      )  ->							Any:
    return np.reshape(__UpperCamelCase       ,     (current.shape[0], current.shape[1] * current.shape[2])      ).T




def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						np.ndarray      )  ->							Any:
    return np.reshape(__UpperCamelCase       ,     (current.shape[0] * current.shape[1], current.shape[2])      ).T




def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						str       ,     __UpperCamelCase			:						List[str]      )  ->							Optional[Any]:
    UpperCAmelCase_							 =     state_dict['''optimizer''']['''target''']['''Transformer'''][f'encoderblock_{i}']

    UpperCAmelCase_							 =     f'encoder.layer.{i}.'
    UpperCAmelCase_							 =     {
        prefix + '''intermediate.dense.bias''': state['''MlpBlock_0''']['''Dense_0''']['''bias'''],
        prefix + '''intermediate.dense.weight''': np.transpose(state['''MlpBlock_0''']['''Dense_0''']['''kernel''']      ),
        prefix + '''output.dense.bias''': state['''MlpBlock_0''']['''Dense_1''']['''bias'''],
        prefix + '''output.dense.weight''': np.transpose(state['''MlpBlock_0''']['''Dense_1''']['''kernel''']      ),
        prefix + '''layernorm_before.bias''': state['''LayerNorm_0''']['''bias'''],
        prefix + '''layernorm_before.weight''': state['''LayerNorm_0''']['''scale'''],
        prefix + '''layernorm_after.bias''': state['''LayerNorm_1''']['''bias'''],
        prefix + '''layernorm_after.weight''': state['''LayerNorm_1''']['''scale'''],
        prefix
        + '''attention.attention.query.bias''': transform_attention(
            state['''MultiHeadDotProductAttention_0''']['''query''']['''bias''']      ),
        prefix
        + '''attention.attention.query.weight''': transform_attention(
            state['''MultiHeadDotProductAttention_0''']['''query''']['''kernel''']      ),
        prefix
        + '''attention.attention.key.bias''': transform_attention(state['''MultiHeadDotProductAttention_0''']['''key''']['''bias''']      ),
        prefix
        + '''attention.attention.key.weight''': transform_attention(
            state['''MultiHeadDotProductAttention_0''']['''key''']['''kernel''']      ),
        prefix
        + '''attention.attention.value.bias''': transform_attention(
            state['''MultiHeadDotProductAttention_0''']['''value''']['''bias''']      ),
        prefix
        + '''attention.attention.value.weight''': transform_attention(
            state['''MultiHeadDotProductAttention_0''']['''value''']['''kernel''']      ),
        prefix + '''attention.output.dense.bias''': state['''MultiHeadDotProductAttention_0''']['''out''']['''bias'''],
        prefix
        + '''attention.output.dense.weight''': transform_attention_output_weight(
            state['''MultiHeadDotProductAttention_0''']['''out''']['''kernel''']      ),
    }

    return new_state




def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						List[str]      )  ->							Optional[int]:
    return sum([1 if '''encoderblock_''' in k else 0 for k in state_dict['''optimizer''']['''target''']['''Transformer'''].keys()]      )




def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						Tuple       ,     __UpperCamelCase			:						Any=False      )  ->							Dict:
    UpperCAmelCase_							 =     get_n_layers(__UpperCamelCase      )

    UpperCAmelCase_							 =     OrderedDict()

    UpperCAmelCase_							 =     state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''bias''']
    UpperCAmelCase_							 =     state_dict['''optimizer''']['''target''']['''Transformer''']['''encoder_norm''']['''scale''']

    UpperCAmelCase_							 =     np.transpose(
        state_dict['''optimizer''']['''target''']['''embedding''']['''kernel''']       ,     (4, 3, 0, 1, 2)      )
    UpperCAmelCase_							 =     state_dict['''optimizer''']['''target''']['''embedding''']['''bias''']

    UpperCAmelCase_							 =     state_dict['''optimizer''']['''target''']['''cls''']
    UpperCAmelCase_							 =     state_dict['''optimizer''']['''target''']['''Transformer''']['''posembed_input'''][
        '''pos_embedding'''
    ]

    for i in range(__UpperCamelCase      ):
        new_state.update(transform_state_encoder_block(__UpperCamelCase       ,     __UpperCamelCase      )      )

    if classification_head:
        UpperCAmelCase_							 =     {'''vivit.''' + k: v for k, v in new_state.items()}
        UpperCAmelCase_							 =     np.transpose(state_dict['''optimizer''']['''target''']['''output_projection''']['''kernel''']      )
        UpperCAmelCase_							 =     np.transpose(state_dict['''optimizer''']['''target''']['''output_projection''']['''bias''']      )

    return {k: torch.tensor(__UpperCamelCase      ) for k, v in new_state.items()}




def 					SCREAMING_SNAKE_CASE	(       )  ->							VivitImageProcessor:
    UpperCAmelCase_							 =     VivitImageProcessor()

    assert extractor.do_resize is True
    assert extractor.size == {"shortest_edge": 256}
    assert extractor.do_center_crop is True
    assert extractor.crop_size == {"width": 224, "height": 224}
    assert extractor.resample == PILImageResampling.BILINEAR

    # here: https://github.com/deepmind/dmvr/blob/master/dmvr/modalities.py
    # one can seen that add_image has default values for normalization_mean and normalization_std set to 0 and 1
    # which effectively means no normalization (and ViViT does not overwrite those when calling this func)
    assert extractor.do_normalize is False
    assert extractor.do_rescale is True
    assert extractor.rescale_factor == 1 / 255

    # zero-centering = True in original implementation
    assert extractor.do_zero_centering is True

    return extractor




def 					SCREAMING_SNAKE_CASE	(       __UpperCamelCase			:						str      )  ->							List[str]:
    UpperCAmelCase_							 =     '''checkpoint'''

    if not os.path.exists(__UpperCamelCase      ):
        download_checkpoint(__UpperCamelCase      )

    UpperCAmelCase_							 =     restore_checkpoint(__UpperCamelCase       ,     __UpperCamelCase      )
    UpperCAmelCase_							 =     transform_state(__UpperCamelCase       ,     classification_head=__UpperCamelCase      )

    UpperCAmelCase_							 =     get_vivit_config()

    assert config.image_size == 224
    assert config.num_frames == 32

    UpperCAmelCase_							 =     VivitForVideoClassification(__UpperCamelCase      )
    model.load_state_dict(__UpperCamelCase      )
    model.eval()

    UpperCAmelCase_							 =     get_processor()

    UpperCAmelCase_							 =     prepare_video()
    UpperCAmelCase_							 =     extractor(__UpperCamelCase       ,     return_tensors='''pt'''      )

    UpperCAmelCase_							 =     model(**__UpperCamelCase      )

    UpperCAmelCase_							 =     torch.Size([1, 400]      )
    UpperCAmelCase_							 =     torch.tensor([-1.0_543, 2.0_764, -0.2_104, 0.4_439, -0.9_658]      )

    assert outputs.logits.shape == expected_shape
    assert torch.allclose(outputs.logits[0, :5]       ,     __UpperCamelCase       ,     atol=1e-4      ), outputs.logits[0, :5]

    model.save_pretrained(__UpperCamelCase      )
    extractor.save_pretrained(__UpperCamelCase      )


if __name__ == "__main__":
 _lowerCamelCase			 =						argparse.ArgumentParser()

 parser.add_argument('--output_model_name', '-o', type=str, help='Output path for the converted HuggingFace model')

 _lowerCamelCase			 =						parser.parse_args()
 convert(args.output_model_name)











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 			a     (   _A			):

     '''simple docstring'''





     lowerCAmelCase				:       List[Any]   		=  'blenderbot'
     lowerCAmelCase				:       int   		=  ['past_key_values']
     lowerCAmelCase				:       int   		=  {'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}
     def __init__(	self   :			int					,    __snake_case   :			Optional[Any]=80_08					,    __snake_case   :			List[str]=1_28					,    __snake_case   :			List[str]=2					,    __snake_case   :			Tuple=1_02_40					,    __snake_case   :			List[Any]=32					,    __snake_case   :			int=24					,    __snake_case   :			str=1_02_40					,    __snake_case   :			Optional[int]=32					,    __snake_case   :			Dict=0.0					,    __snake_case   :			int=0.0					,    __snake_case   :			str=True					,    __snake_case   :			Dict=True					,    __snake_case   :			Dict="gelu"					,    __snake_case   :			List[str]=25_60					,    __snake_case   :			List[Any]=0.1					,    __snake_case   :			str=0.0					,    __snake_case   :			Dict=0.0					,    __snake_case   :			Tuple=0.02					,    __snake_case   :			Any=1					,    __snake_case   :			int=False					,    __snake_case   :			Any=0					,    __snake_case   :			Optional[Any]=1					,    __snake_case   :			Optional[int]=2					,    __snake_case   :			List[Any]=3					,    __snake_case   :			int=2					,    **__snake_case   :			Tuple					,    ):
         UpperCAmelCase_							 =     vocab_size
         UpperCAmelCase_							 =     max_position_embeddings
         UpperCAmelCase_							 =     d_model
         UpperCAmelCase_							 =     encoder_ffn_dim
         UpperCAmelCase_							 =     encoder_layers
         UpperCAmelCase_							 =     encoder_attention_heads
         UpperCAmelCase_							 =     decoder_ffn_dim
         UpperCAmelCase_							 =     decoder_layers
         UpperCAmelCase_							 =     decoder_attention_heads
         UpperCAmelCase_							 =     dropout
         UpperCAmelCase_							 =     attention_dropout
         UpperCAmelCase_							 =     activation_dropout
         UpperCAmelCase_							 =     activation_function
         UpperCAmelCase_							 =     init_std
         UpperCAmelCase_							 =     encoder_layerdrop
         UpperCAmelCase_							 =     decoder_layerdrop
         UpperCAmelCase_							 =     use_cache
         UpperCAmelCase_							 =     encoder_layers
         UpperCAmelCase_							 =     scale_embedding  # scale factor will be sqrt(d_model) if True

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



class 			a     (   _A			):

     '''simple docstring'''





     @property
     def   lowerCamelCase_   (	self   :			Any				):
         if self.task in ["default", "seq2seq-lm"]:
             UpperCAmelCase_							 =     OrderedDict(
                 [
                     ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
                     ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
                 ]				)
             if self.use_past:
                 UpperCAmelCase_							 =     {0: '''batch'''}
                 UpperCAmelCase_							 =     {0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
             else:
                 UpperCAmelCase_							 =     {0: '''batch''', 1: '''decoder_sequence'''}
                 UpperCAmelCase_							 =     {0: '''batch''', 1: '''decoder_sequence'''}
             if self.use_past:
                 self.fill_with_past_key_values_(__snake_case					,    direction='''inputs'''				)
         elif self.task == "causal-lm":
             UpperCAmelCase_							 =     OrderedDict(
                 [
                     ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
                     ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
                 ]				)
             if self.use_past:
                 UpperCAmelCase_ ,							UpperCAmelCase_							 =     self.num_layers
                 for i in range(__snake_case				):
                     UpperCAmelCase_							 =     {0: '''batch''', 2: '''past_sequence + sequence'''}
                     UpperCAmelCase_							 =     {0: '''batch''', 2: '''past_sequence + sequence'''}
         else:
             UpperCAmelCase_							 =     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   lowerCamelCase_   (	self   :			Optional[Any]				):
         if self.task in ["default", "seq2seq-lm"]:
             UpperCAmelCase_							 =     super().outputs
         else:
             UpperCAmelCase_							 =     super(__snake_case					,    self				).outputs
             if self.use_past:
                 UpperCAmelCase_ ,							UpperCAmelCase_							 =     self.num_layers
                 for i in range(__snake_case				):
                     UpperCAmelCase_							 =     {0: '''batch''', 2: '''past_sequence + sequence'''}
                     UpperCAmelCase_							 =     {0: '''batch''', 2: '''past_sequence + sequence'''}
         return common_outputs
     def   lowerCamelCase_   (	self   :			Optional[Any]					,    __snake_case   :			PreTrainedTokenizer					,    __snake_case   :			int = -1					,    __snake_case   :			int = -1					,    __snake_case   :			bool = False					,    __snake_case   :			Optional[TensorType] = None					,    ):
         UpperCAmelCase_							 =     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             __snake_case					,    __snake_case					,    __snake_case					,    __snake_case					,    __snake_case				)
         # Generate decoder inputs
         UpperCAmelCase_							 =     seq_length if not self.use_past else 1
         UpperCAmelCase_							 =     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             __snake_case					,    __snake_case					,    __snake_case					,    __snake_case					,    __snake_case				)
         UpperCAmelCase_							 =     {F'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
         UpperCAmelCase_							 =     dict(**__snake_case					,    **__snake_case				)

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

             for _ in range(__snake_case				):
                 common_inputs["past_key_values"].append(
                     (
                         torch.zeros(__snake_case				),
                         torch.zeros(__snake_case				),
                         torch.zeros(__snake_case				),
                         torch.zeros(__snake_case				),
                     )				)
         return common_inputs
     def   lowerCamelCase_   (	self   :			Optional[Any]					,    __snake_case   :			PreTrainedTokenizer					,    __snake_case   :			int = -1					,    __snake_case   :			int = -1					,    __snake_case   :			bool = False					,    __snake_case   :			Optional[TensorType] = None					,    ):
         UpperCAmelCase_							 =     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             __snake_case					,    __snake_case					,    __snake_case					,    __snake_case					,    __snake_case				)

         if self.use_past:
             if not is_torch_available():
                 raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''				)
             else:
                 import torch
             UpperCAmelCase_ ,							UpperCAmelCase_							 =     common_inputs['''input_ids'''].shape
             UpperCAmelCase_							 =     seqlen
             UpperCAmelCase_ ,							UpperCAmelCase_							 =     self.num_layers
             UpperCAmelCase_ ,							UpperCAmelCase_							 =     self.num_attention_heads
             UpperCAmelCase_							 =     (
                 batch,
                 num_encoder_attention_heads,
                 past_key_values_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )
             UpperCAmelCase_							 =     common_inputs['''attention_mask'''].dtype
             UpperCAmelCase_							 =     torch.cat(
                 [common_inputs['''attention_mask'''], torch.ones(__snake_case					,    __snake_case					,    dtype=__snake_case				)]					,    dim=1				)
             UpperCAmelCase_							 =     [
                 (torch.zeros(__snake_case				), torch.zeros(__snake_case				)) for _ in range(__snake_case				)
             ]
         return common_inputs
     def   lowerCamelCase_   (	self   :			Any					,    __snake_case   :			PreTrainedTokenizer					,    __snake_case   :			int = -1					,    __snake_case   :			int = -1					,    __snake_case   :			bool = False					,    __snake_case   :			Optional[TensorType] = None					,    ):
         # 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
         UpperCAmelCase_							 =     compute_effective_axis_dimension(
             __snake_case					,    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
         UpperCAmelCase_							 =     tokenizer.num_special_tokens_to_add(__snake_case				)
         UpperCAmelCase_							 =     compute_effective_axis_dimension(
             __snake_case					,    fixed_dimension=OnnxConfig.default_fixed_sequence					,    num_token_to_add=__snake_case				)

         # Generate dummy inputs according to compute batch and sequence
         UpperCAmelCase_							 =     [''' '''.join([tokenizer.unk_token]				) * seq_length] * batch_size
         UpperCAmelCase_							 =     dict(tokenizer(__snake_case					,    return_tensors=__snake_case				)				)
         return common_inputs
     def   lowerCamelCase_   (	self   :			List[Any]					,    __snake_case   :			PreTrainedTokenizer					,    __snake_case   :			int = -1					,    __snake_case   :			int = -1					,    __snake_case   :			bool = False					,    __snake_case   :			Optional[TensorType] = None					,    ):
         if self.task in ["default", "seq2seq-lm"]:
             UpperCAmelCase_							 =     self._generate_dummy_inputs_for_default_and_seqaseq_lm(
                 __snake_case					,    batch_size=__snake_case					,    seq_length=__snake_case					,    is_pair=__snake_case					,    framework=__snake_case				)

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

         return common_inputs
     def   lowerCamelCase_   (	self   :			Dict					,    __snake_case   :			Optional[int]					,    __snake_case   :			int					,    __snake_case   :			str					,    __snake_case   :			Optional[int]				):
         if self.task in ["default", "seq2seq-lm"]:
             UpperCAmelCase_							 =     super()._flatten_past_key_values_(__snake_case					,    __snake_case					,    __snake_case					,    __snake_case				)
         else:
             UpperCAmelCase_							 =     super(__snake_case					,    self				)._flatten_past_key_values_(
                 __snake_case					,    __snake_case					,    __snake_case					,    __snake_case				)




     def   lowerCamelCase_   (	self   :			Dict					,    __snake_case   :			Mapping[str, Mapping[int, str]]					,    __snake_case   :			str				):
         if direction not in ["inputs", "outputs"]:
             raise ValueError(F'direction must either be "inputs" or "outputs", but {direction} was given'				)

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

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

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





