
'''simple docstring'''
import argparse
import collections

import jax
import jax.numpy as jnp
import torch
import torch.nn as nn
from clip.model import CLIP
from flax.training import checkpoints
from huggingface_hub import Repository

from transformers import (
    CLIPTokenizer,
    OwlViTConfig,
    OwlViTForObjectDetection,
    OwlViTImageProcessor,
    OwlViTModel,
    OwlViTProcessor,
)


__lowerCAmelCase  							= {
    'vit_b32': {
        'embed_dim': 512,
        'image_resolution': 768,
        'context_length': 16,
        'vocab_size': 49_408,
        'vision_layers': 12,
        'vision_width': 768,
        'vision_patch_size': 32,
        'transformer_width': 512,
        'transformer_heads': 8,
        'transformer_layers': 12,
    },
    'vit_b16': {
        'embed_dim': 512,
        'image_resolution': 768,
        'context_length': 16,
        'vocab_size': 49_408,
        'vision_layers': 12,
        'vision_width': 768,
        'vision_patch_size': 16,
        'transformer_width': 512,
        'transformer_heads': 8,
        'transformer_layers': 12,
    },
    'vit_l14': {
        'embed_dim': 768,
        'image_resolution': 840,
        'context_length': 16,
        'vocab_size': 49_408,
        'vision_layers': 24,
        'vision_width': 1_024,
        'vision_patch_size': 14,
        'transformer_width': 768,
        'transformer_heads': 12,
        'transformer_layers': 12,
    },
}






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE=""							,      _SCREAMING_SNAKE_CASE="/"   ):
							_snake_case     =						[]

							for k, v in params.items():
														_snake_case     =						parent_key + sep + k if parent_key else k

														if isinstance(_SCREAMING_SNAKE_CASE							,      collections.MutableMapping   ):
																					items.extend(flatten_nested_dict(_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE							,      sep=_SCREAMING_SNAKE_CASE   ).items()   )
														else:
																					items.append((new_key, v)   )
							return dict(_SCREAMING_SNAKE_CASE   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE   ):
							return jax.tree_util.tree_map(lambda _SCREAMING_SNAKE_CASE   : x.astype(jnp.floataa   ) if x.dtype == jnp.bfloataa else x							,      _SCREAMING_SNAKE_CASE   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							_snake_case,							_snake_case,							_snake_case     =						pt_attn_layer.in_proj_weight.chunk(3							,      dim=0   )
							_snake_case,							_snake_case,							_snake_case     =						pt_attn_layer.in_proj_bias.chunk(3							,      dim=0   )

							_snake_case     =						pt_attn_layer.out_proj.weight
							_snake_case     =						pt_attn_layer.out_proj.bias

							_snake_case     =						q_proj
							_snake_case     =						q_proj_bias

							_snake_case     =						k_proj
							_snake_case     =						k_proj_bias

							_snake_case     =						v_proj
							_snake_case     =						v_proj_bias

							_snake_case     =						out_proj_weights
							_snake_case     =						out_proj_bias






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							copy_linear(hf_mlp.fca							,      pt_mlp.c_fc   )
							copy_linear(hf_mlp.fca							,      pt_mlp.c_proj   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							_snake_case     =						pt_linear.weight
							_snake_case     =						pt_linear.bias






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							# copy layer norms
							copy_linear(hf_layer.layer_norma							,      pt_layer.ln_a   )
							copy_linear(hf_layer.layer_norma							,      pt_layer.ln_a   )

							# copy MLP
							copy_mlp(hf_layer.mlp							,      pt_layer.mlp   )

							# copy attn
							copy_attn_layer(hf_layer.self_attn							,      pt_layer.attn   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							for hf_layer, pt_layer in zip(_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
														copy_layer(_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							# copy  embeds
							_snake_case     =						pt_model.token_embedding.weight
							_snake_case     =						pt_model.positional_embedding

							# copy layer norm
							copy_linear(hf_encoder.final_layer_norm							,      pt_model.ln_final   )

							# copy hidden layers
							copy_layers(hf_encoder.encoder.layers							,      pt_model.transformer.resblocks   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							# copy projection
							_snake_case     =						pt_model.text_projection.data.T

							# copy text encoder
							copy_encoder(hf_model.text_model							,      _SCREAMING_SNAKE_CASE   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							# copy projection
							_snake_case     =						pt_model.visual.proj.data.T

							# copy layer norms
							copy_linear(hf_model.vision_model.pre_layernorm							,      pt_model.visual.ln_pre   )
							copy_linear(hf_model.vision_model.post_layernorm							,      pt_model.visual.ln_post   )

							# copy embeds
							_snake_case     =						pt_model.visual.conva.weight.data
							_snake_case     =						pt_model.visual.class_embedding
							_snake_case     =						pt_model.visual.positional_embedding.data

							# copy encoder
							copy_layers(hf_model.vision_model.encoder.layers							,      pt_model.visual.transformer.resblocks   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							_snake_case     =						flatten_nested_dict(flax_params["""backbone"""]["""merged_class_token"""]   )

							_snake_case     =						torch.from_numpy(flax_class_token_params["""scale"""]   )
							_snake_case     =						torch.from_numpy(flax_class_token_params["""bias"""]   )
							_snake_case     =						nn.Parameter(_SCREAMING_SNAKE_CASE   )
							_snake_case     =						nn.Parameter(_SCREAMING_SNAKE_CASE   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							_snake_case     =						hf_model.state_dict()
							_snake_case     =						{}

							# Rename class prediction head flax params to pytorch HF
							_snake_case     =						flatten_nested_dict(flax_params["""class_head"""]   )

							for flax_key, v in flax_class_params.items():
														_snake_case     =						flax_key.replace("""/"""							,      """."""   )
														_snake_case     =						torch_key.replace(""".kernel"""							,      """.weight"""   )
														_snake_case     =						torch_key.replace("""Dense_0"""							,      """dense0"""   )
														_snake_case     =						"""class_head.""" + torch_key

														if "weight" in torch_key and v.ndim == 2:
																					_snake_case     =						v.T

														_snake_case     =						nn.Parameter(torch.from_numpy(_SCREAMING_SNAKE_CASE   )   )

							# Rename box prediction box flax params to pytorch HF
							_snake_case     =						flatten_nested_dict(flax_params["""obj_box_head"""]   )

							for flax_key, v in flax_box_params.items():
														_snake_case     =						flax_key.replace("""/"""							,      """."""   )
														_snake_case     =						torch_key.replace(""".kernel"""							,      """.weight"""   )
														_snake_case     =						torch_key.replace("""_"""							,      """"""   ).lower()
														_snake_case     =						"""box_head.""" + torch_key

														if "weight" in torch_key and v.ndim == 2:
																					_snake_case     =						v.T

														_snake_case     =						nn.Parameter(torch.from_numpy(_SCREAMING_SNAKE_CASE   )   )

							# Copy flax params to PyTorch params
							for name, param in new_params.items():
														if name in pt_params.keys():
																					pt_params[name].copy_(_SCREAMING_SNAKE_CASE   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							for k, v in flax_attn_params.items():
														if k.startswith("""transformer"""   ):
																					_snake_case     =						k.replace("""transformer.resblocks"""							,      """text_model.encoder.layers"""   )
														else:
																					_snake_case     =						k.replace("""visual.transformer.resblocks"""							,      """vision_model.encoder.layers"""   )

														_snake_case     =						torch_key.replace("""attn"""							,      """self_attn"""   )
														_snake_case     =						torch_key.replace("""key"""							,      """k_proj"""   )
														_snake_case     =						torch_key.replace("""value"""							,      """v_proj"""   )
														_snake_case     =						torch_key.replace("""query"""							,      """q_proj"""   )
														_snake_case     =						torch_key.replace("""out"""							,      """out_proj"""   )

														if "bias" in torch_key and v.ndim == 2:
																					_snake_case     =						v.shape[0] * v.shape[1]
																					_snake_case     =						v.reshape(_SCREAMING_SNAKE_CASE   )

														if "weight" in torch_key and "out" in torch_key:
																					_snake_case     =						(v.shape[0] * v.shape[1], v.shape[2])
																					_snake_case     =						v.reshape(_SCREAMING_SNAKE_CASE   ).T

														if "weight" in torch_key and "out" not in torch_key:
																					_snake_case     =						(v.shape[0], v.shape[1] * v.shape[2])
																					_snake_case     =						v.reshape(_SCREAMING_SNAKE_CASE   ).T

														# Copy flax CLIP attn params to HF PyTorch params
														_snake_case     =						torch.from_numpy(_SCREAMING_SNAKE_CASE   )
														hf_backbone.state_dict()[torch_key].copy_(_SCREAMING_SNAKE_CASE   )






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE   ):
							_snake_case     =						{}
							_snake_case     =						[]

							for k, v in params.items():
														if "attn." in k:
																					_snake_case     =						k[: k.rindex("""attn."""   ) + 5]
																					if base in processed_attn_layers:
																												continue

																					processed_attn_layers.append(_SCREAMING_SNAKE_CASE   )
																					_snake_case     =						params[base + """out.weight"""].shape[-1]
																					_snake_case     =						params[base + """out.weight"""].reshape(_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ).T
																					_snake_case     =						params[base + """out.bias"""]
														else:
																					_snake_case     =						v
							return new_params






def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   ):
							_snake_case     =						CLIP(**_SCREAMING_SNAKE_CASE   )
							torch_model.eval()
							_snake_case     =						torch_model.state_dict()

							_snake_case     =						flatten_nested_dict(flax_params["""backbone"""]["""clip"""]   )
							_snake_case     =						{}

							for flax_key, v in flax_clip_params.items():
														_snake_case     =						flax_key.replace("""/"""							,      """."""   )
														_snake_case     =						torch_key.replace("""text.token_embedding.embedding"""							,      """token_embedding.kernel"""   )

														if (
														    torch_key.startswith("""text.transformer"""   )
														    or torch_key.startswith("""text.text_projection"""   )
														    or torch_key.startswith("""text.ln_final"""   )
														    or torch_key.startswith("""text.positional_embedding"""   )
														):
																					_snake_case     =						torch_key[5:]

														_snake_case     =						torch_key.replace("""text_projection.kernel"""							,      """text_projection"""   )
														_snake_case     =						torch_key.replace("""visual.proj.kernel"""							,      """visual.proj"""   )
														_snake_case     =						torch_key.replace(""".scale"""							,      """.weight"""   )
														_snake_case     =						torch_key.replace(""".kernel"""							,      """.weight"""   )

														if "conv" in torch_key or "downsample.0.weight" in torch_key:
																					_snake_case     =						v.transpose(3							,      2							,      0							,      1   )

														elif "weight" in torch_key and v.ndim == 2 and "embedding" not in torch_key:
																					# Fully connected layers are transposed, embeddings are not
																					_snake_case     =						v.T

														_snake_case     =						v

							_snake_case     =						_convert_attn_layers(_SCREAMING_SNAKE_CASE   )
							new_torch_params.update(_SCREAMING_SNAKE_CASE   )
							_snake_case     =						{}

							# Copy flax CLIP backbone params to PyTorch params
							for name, param in new_torch_params.items():
														if name in torch_clip_params.keys():
																					_snake_case     =						torch.from_numpy(new_torch_params[name]   )
																					torch_clip_params[name].copy_(_SCREAMING_SNAKE_CASE   )
														else:
																					_snake_case     =						param

							return torch_clip_params, torch_model, attn_params






@torch.no_grad()
def       __SCREAMING_SNAKE_CASE   (						_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE=None   ):
							_snake_case     =						Repository(_SCREAMING_SNAKE_CASE							,      clone_from=f"""google/{pytorch_dump_folder_path}"""   )
							repo.git_pull()

							if config_path is not None:
														_snake_case     =						OwlViTConfig.from_pretrained(_SCREAMING_SNAKE_CASE   )
							else:
														_snake_case     =						OwlViTConfig()

							_snake_case     =						OwlViTModel(_SCREAMING_SNAKE_CASE   ).eval()
							_snake_case     =						OwlViTForObjectDetection(_SCREAMING_SNAKE_CASE   ).eval()

							copy_text_model_and_projection(_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   )
							copy_vision_model_and_projection(_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   )
							_snake_case     =						pt_backbone.logit_scale
							copy_flax_attn_params(_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   )

							_snake_case     =						hf_backbone
							copy_class_merge_token(_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   )
							copy_class_box_heads(_SCREAMING_SNAKE_CASE							,      _SCREAMING_SNAKE_CASE   )

							# Save HF model
							hf_model.save_pretrained(repo.local_dir   )

							# Initialize image processor
							_snake_case     =						OwlViTImageProcessor(
							    size=config.vision_config.image_size							,      crop_size=config.vision_config.image_size   )
							# Initialize tokenizer
							_snake_case     =						CLIPTokenizer.from_pretrained("""openai/clip-vit-base-patch32"""							,      pad_token="""!"""							,      model_max_length=16   )

							# Initialize processor
							_snake_case     =						OwlViTProcessor(image_processor=_SCREAMING_SNAKE_CASE							,      tokenizer=_SCREAMING_SNAKE_CASE   )
							image_processor.save_pretrained(repo.local_dir   )
							processor.save_pretrained(repo.local_dir   )

							repo.git_add()
							repo.git_commit("""Upload model and processor"""   )
							repo.git_push()


if __name__ == "__main__":
					__lowerCAmelCase  							= argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '--owlvit_version',
					    default=None,
					    type=str,
					    required=True,
					    help='OWL-ViT model name [clip_b16, clip_b32, clip_l14].',
					)
					parser.add_argument(
					    '--owlvit_checkpoint', default=None, type=str, required=True, help='Path to flax model checkpoint.'
					)
					parser.add_argument('--hf_config', default=None, type=str, required=True, help='Path to HF model config.')
					parser.add_argument(
					    '--pytorch_dump_folder_path', default='hf_model', type=str, help='Path to the output PyTorch model.'
					)
					__lowerCAmelCase  							= parser.parse_args()

					# Initialize PyToch clip model
					__lowerCAmelCase  							= args.owlvit_version
					if model_name == "clip_b16":
										__lowerCAmelCase  							= CONFIGS['vit_b16']
					elif model_name == "clip_b32":
										__lowerCAmelCase  							= CONFIGS['vit_b32']
					elif model_name == "clip_l14":
										__lowerCAmelCase  							= CONFIGS['vit_l14']

					# Load from checkpoint and convert params to float-32
					__lowerCAmelCase  							= checkpoints.restore_checkpoint(args.owlvit_checkpoint, target=None)['optimizer']['target']
					__lowerCAmelCase  							= jax.tree_util.tree_map(lambda x: x.astype(jnp.floataa) if x.dtype == jnp.bfloataa else x, variables)
					del variables

					# Convert CLIP backbone
					__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase  							= convert_clip_backbone(flax_params, torch_config)

					convert_owlvit_checkpoint(clip_pt, flax_params, attn_params, args.pytorch_dump_folder_path, args.hf_config)

'''simple docstring'''
from collections import OrderedDict
from typing import Any, Mapping, Optional

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


__lowerCAmelCase  							= logging.get_logger(__name__)

__lowerCAmelCase  							= {
    'facebook/m2m100_418M': 'https://huggingface.co/facebook/m2m100_418M/resolve/main/config.json',
    # See all M2M100 models at https://huggingface.co/models?filter=m2m_100
}


class     _lowerCAmelCase					(					__snake_case   ):







							'''simple docstring'''

							lowerCAmelCase_      		=		"m2m_100"
							lowerCAmelCase_      		=		["past_key_values"]
							lowerCAmelCase_      		=		{"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}

							def __init__(self						,						UpperCAmelCase=128112						,						UpperCAmelCase=1024						,						UpperCAmelCase=12						,						UpperCAmelCase=4096						,						UpperCAmelCase=16						,						UpperCAmelCase=12						,						UpperCAmelCase=4096						,						UpperCAmelCase=16						,						UpperCAmelCase=0.05						,						UpperCAmelCase=0.05						,						UpperCAmelCase=True						,						UpperCAmelCase=True						,						UpperCAmelCase="relu"						,						UpperCAmelCase=1024						,						UpperCAmelCase=0.1						,						UpperCAmelCase=0.1						,						UpperCAmelCase=0.0						,						UpperCAmelCase=0.02						,						UpperCAmelCase=2						,						UpperCAmelCase=True						,						UpperCAmelCase=1						,						UpperCAmelCase=0						,						UpperCAmelCase=2						,						**UpperCAmelCase						,						)  ->     int:
														_snake_case     =						vocab_size
														_snake_case     =						max_position_embeddings
														_snake_case     =						d_model
														_snake_case     =						encoder_ffn_dim
														_snake_case     =						encoder_layers
														_snake_case     =						encoder_attention_heads
														_snake_case     =						decoder_ffn_dim
														_snake_case     =						decoder_layers
														_snake_case     =						decoder_attention_heads
														_snake_case     =						dropout
														_snake_case     =						attention_dropout
														_snake_case     =						activation_dropout
														_snake_case     =						activation_function
														_snake_case     =						init_std
														_snake_case     =						encoder_layerdrop
														_snake_case     =						decoder_layerdrop
														_snake_case     =						use_cache
														_snake_case     =						encoder_layers
														_snake_case     =						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						,						**UpperCAmelCase						,						)


class     _lowerCAmelCase					(					__snake_case   ):







							'''simple docstring'''

							@property
							def  lowercase							(self						)  ->     Mapping[str, Mapping[int, str]]:
														_snake_case     =						OrderedDict(
														    [
														        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
														        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
														    ]						)

														if self.use_past:
																					_snake_case     =						{0: """batch"""}
																					_snake_case     =						{0: """batch""", 1: """past_decoder_sequence + sequence"""}
														else:
																					_snake_case     =						{0: """batch""", 1: """decoder_sequence"""}
																					_snake_case     =						{0: """batch""", 1: """decoder_sequence"""}

														if self.use_past:
																					self.fill_with_past_key_values_(UpperCAmelCase						,						direction="""inputs"""						)
														return common_inputs

							def  lowercase							(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
														_snake_case     =						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
														_snake_case     =						tokenizer.num_special_tokens_to_add(UpperCAmelCase						)
														_snake_case     =						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
														_snake_case     =						[""" """.join([tokenizer.unk_token]						) * seq_length] * batch_size
														_snake_case     =						dict(tokenizer(UpperCAmelCase						,						return_tensors=UpperCAmelCase						)						)
														return common_inputs

							def  lowercase							(self						,						UpperCAmelCase						,						UpperCAmelCase = -1						,						UpperCAmelCase = -1						,						UpperCAmelCase = False						,						UpperCAmelCase = None						,						)  ->     Mapping[str, Any]:
														_snake_case     =						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
														    UpperCAmelCase						,						UpperCAmelCase						,						UpperCAmelCase						,						UpperCAmelCase						,						UpperCAmelCase						)

														# Generate decoder inputs
														_snake_case     =						seq_length if not self.use_past else 1
														_snake_case     =						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
														    UpperCAmelCase						,						UpperCAmelCase						,						UpperCAmelCase						,						UpperCAmelCase						,						UpperCAmelCase						)
														_snake_case     =						{f"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
														_snake_case     =						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
																					_snake_case,							_snake_case     =						common_inputs["""input_ids"""].shape
																					_snake_case     =						common_inputs["""decoder_input_ids"""].shape[1]
																					_snake_case,							_snake_case     =						self.num_attention_heads
																					_snake_case     =						(
																					    batch,
																					    num_encoder_attention_heads,
																					    encoder_seq_length,
																					    self._config.hidden_size // num_encoder_attention_heads,
																					)
																					_snake_case     =						decoder_seq_length + 3
																					_snake_case     =						(
																					    batch,
																					    num_decoder_attention_heads,
																					    decoder_past_length,
																					    self._config.hidden_size // num_decoder_attention_heads,
																					)

																					_snake_case     =						torch.cat(
																					    [common_inputs["""decoder_attention_mask"""], torch.ones(UpperCAmelCase						,						UpperCAmelCase						)]						,						dim=1						)

																					_snake_case     =						[]
																					# If the number of encoder and decoder layers are present in the model configuration, both are considered
																					_snake_case,							_snake_case     =						self.num_layers
																					_snake_case     =						min(UpperCAmelCase						,						UpperCAmelCase						)
																					_snake_case     =						max(UpperCAmelCase						,						UpperCAmelCase						) - min_num_layers
																					_snake_case     =						"""encoder""" if num_encoder_layers > num_decoder_layers else """decoder"""

																					for _ in range(UpperCAmelCase						):
																												common_inputs["past_key_values"].append(
																												    (
																												        torch.zeros(UpperCAmelCase						),
																												        torch.zeros(UpperCAmelCase						),
																												        torch.zeros(UpperCAmelCase						),
																												        torch.zeros(UpperCAmelCase						),
																												    )						)
																					# TODO: test this.
																					_snake_case     =						encoder_shape if remaining_side_name == """encoder""" else decoder_shape
																					for _ in range(UpperCAmelCase						,						UpperCAmelCase						):
																												common_inputs["past_key_values"].append((torch.zeros(UpperCAmelCase						), torch.zeros(UpperCAmelCase						))						)
														return common_inputs





							lowerCAmelCase_      		=		_generate_dummy_inputs_for_default_and_seqaseq_lm