


'''simple docstring'''

import collections
import inspect
import unittest

from transformers import FocalNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import (
     FocalNetBackbone,
     FocalNetForImageClassification,
     FocalNetForMaskedImageModeling,
     FocalNetModel,
 )
 from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor
class     __lowerCAmelCase		:





      '''simple docstring'''






      def __init__(self       :       List[Any]  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       int=13  ,		UpperCamelCase       :       Any=32  ,		UpperCamelCase       :       Union[str, Any]=2  ,		UpperCamelCase       :       Tuple=3  ,		UpperCamelCase       :       str=16  ,		UpperCamelCase       :       int=[32, 64, 128]  ,		UpperCamelCase       :       Union[str, Any]=[1, 2, 1]  ,		UpperCamelCase       :       Optional[Any]=[2, 2, 4]  ,		UpperCamelCase       :       str=2  ,		UpperCamelCase       :       str=2.0  ,		UpperCamelCase       :       Tuple=True  ,		UpperCamelCase       :       Union[str, Any]=0.0  ,		UpperCamelCase       :       Union[str, Any]=0.0  ,		UpperCamelCase       :       Optional[Any]=0.1  ,		UpperCamelCase       :       Tuple="gelu"  ,		UpperCamelCase       :       int=False  ,		UpperCamelCase       :       int=True  ,		UpperCamelCase       :       List[Any]=0.02  ,		UpperCamelCase       :       Dict=1E-5  ,		UpperCamelCase       :       Tuple=True  ,		UpperCamelCase       :       Dict=None  ,		UpperCamelCase       :       Any=True  ,		UpperCamelCase       :       str=10  ,		UpperCamelCase       :       List[str]=8  ,		UpperCamelCase       :       Optional[int]=["stage1", "stage2"]  ,		UpperCamelCase       :       Any=[1, 2]  ,		):

             '''simple docstring'''







             lowercase__      =					parent
             lowercase__      =					batch_size
             lowercase__      =					image_size
             lowercase__      =					patch_size
             lowercase__      =					num_channels
             lowercase__      =					embed_dim
             lowercase__      =					hidden_sizes
             lowercase__      =					depths
             lowercase__      =					num_heads
             lowercase__      =					window_size
             lowercase__      =					mlp_ratio
             lowercase__      =					qkv_bias
             lowercase__      =					hidden_dropout_prob
             lowercase__      =					attention_probs_dropout_prob
             lowercase__      =					drop_path_rate
             lowercase__      =					hidden_act
             lowercase__      =					use_absolute_embeddings
             lowercase__      =					patch_norm
             lowercase__      =					layer_norm_eps
             lowercase__      =					initializer_range
             lowercase__      =					is_training
             lowercase__      =					scope
             lowercase__      =					use_labels
             lowercase__      =					type_sequence_label_size
             lowercase__      =					encoder_stride
             lowercase__      =					out_features
             lowercase__      =					out_indices



      def  UpperCamelCase__						(self       :       List[Any]					):

             '''simple docstring'''







             lowercase__      =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

             lowercase__      =					None
             if self.use_labels:
                    lowercase__      =					ids_tensor([self.batch_size]  ,		self.type_sequence_label_size					)

             lowercase__      =					self.get_config()

             return config, pixel_values, labels



      def  UpperCamelCase__						(self       :       Any					):

             '''simple docstring'''







             return FocalNetConfig(
                 image_size=self.image_size  ,		patch_size=self.patch_size  ,		num_channels=self.num_channels  ,		embed_dim=self.embed_dim  ,		hidden_sizes=self.hidden_sizes  ,		depths=self.depths  ,		num_heads=self.num_heads  ,		window_size=self.window_size  ,		mlp_ratio=self.mlp_ratio  ,		qkv_bias=self.qkv_bias  ,		hidden_dropout_prob=self.hidden_dropout_prob  ,		attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,		drop_path_rate=self.drop_path_rate  ,		hidden_act=self.hidden_act  ,		use_absolute_embeddings=self.use_absolute_embeddings  ,		path_norm=self.patch_norm  ,		layer_norm_eps=self.layer_norm_eps  ,		initializer_range=self.initializer_range  ,		encoder_stride=self.encoder_stride  ,		out_features=self.out_features  ,		out_indices=self.out_indices  ,		)



      def  UpperCamelCase__						(self       :       List[str]  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       Optional[Any]  ,		UpperCamelCase       :       Optional[int]					):

             '''simple docstring'''







             lowercase__      =					FocalNetModel(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase					)

             lowercase__      =					((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths					) - 1))
             lowercase__      =					int(config.embed_dim * 2 ** (len(config.depths					) - 1)					)

             self.parent.assertEqual(result.last_hidden_state.shape  ,		(self.batch_size, expected_seq_len, expected_dim)					)



      def  UpperCamelCase__						(self       :       Any  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       List[str]  ,		UpperCamelCase       :       Any					):

             '''simple docstring'''







             lowercase__      =					FocalNetBackbone(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase					)

             # verify feature maps
             self.parent.assertEqual(len(result.feature_maps					)  ,		len(config.out_features					)					)
             self.parent.assertListEqual(list(result.feature_maps[0].shape					)  ,		[self.batch_size, self.image_size, 8, 8]					)

             # verify channels
             self.parent.assertEqual(len(model.channels					)  ,		len(config.out_features					)					)
             self.parent.assertListEqual(model.channels  ,		config.hidden_sizes[:-1]					)

             # verify backbone works with out_features=None
             lowercase__      =					None
             lowercase__      =					FocalNetBackbone(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase					)

             # verify feature maps
             self.parent.assertEqual(len(result.feature_maps					)  ,		1					)
             self.parent.assertListEqual(list(result.feature_maps[0].shape					)  ,		[self.batch_size, self.image_size * 2, 4, 4]					)

             # verify channels
             self.parent.assertEqual(len(model.channels					)  ,		1					)
             self.parent.assertListEqual(model.channels  ,		[config.hidden_sizes[-1]]					)



      def  UpperCamelCase__						(self       :       Union[str, Any]  ,		UpperCamelCase       :       List[str]  ,		UpperCamelCase       :       Union[str, Any]  ,		UpperCamelCase       :       str					):

             '''simple docstring'''







             lowercase__      =					FocalNetForMaskedImageModeling(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase					)
             self.parent.assertEqual(
                 result.reconstruction.shape  ,		(self.batch_size, self.num_channels, self.image_size, self.image_size)					)

             # test greyscale images
             lowercase__      =					1
             lowercase__      =					FocalNetForMaskedImageModeling(UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()

             lowercase__      =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
             lowercase__      =					model(UpperCamelCase					)
             self.parent.assertEqual(result.reconstruction.shape  ,		(self.batch_size, 1, self.image_size, self.image_size)					)



      def  UpperCamelCase__						(self       :       Dict  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       Optional[Any]					):

             '''simple docstring'''







             lowercase__      =					self.type_sequence_label_size
             lowercase__      =					FocalNetForImageClassification(UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase  ,		labels=UpperCamelCase					)
             self.parent.assertEqual(result.logits.shape  ,		(self.batch_size, self.type_sequence_label_size)					)

             # test greyscale images
             lowercase__      =					1
             lowercase__      =					FocalNetForImageClassification(UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()

             lowercase__      =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
             lowercase__      =					model(UpperCamelCase					)
             self.parent.assertEqual(result.logits.shape  ,		(self.batch_size, self.type_sequence_label_size)					)






      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             lowercase__      =					self.prepare_config_and_inputs()

             lowercase__     ,lowercase__     ,lowercase__      =					config_and_inputs
             lowercase__      =					{'''pixel_values''': pixel_values}
             return config, inputs_dict






@require_torch
class     __lowerCAmelCase		(lowercase_    ,	lowercase_    ,	unittest.TestCase	):





      '''simple docstring'''






      lowerCAmelCase__   :	Any       =				(
          (
              FocalNetModel,
              FocalNetForImageClassification,
              FocalNetForMaskedImageModeling,
              FocalNetBackbone,
          )
          if is_torch_available()
          else ()
      )
      lowerCAmelCase__   :	Any       =				(
          {"""feature-extraction""": FocalNetModel, """image-classification""": FocalNetForImageClassification}
          if is_torch_available()
          else {}
      )
      lowerCAmelCase__   :	List[str]       =				False

      lowerCAmelCase__   :	Any       =				False
      lowerCAmelCase__   :	Optional[Any]       =				False
      lowerCAmelCase__   :	List[Any]       =				False
      lowerCAmelCase__   :	Tuple       =				False



      def  UpperCamelCase__						(self       :       int					):

             '''simple docstring'''







             lowercase__      =					FocalNetModelTester(self					)
             lowercase__      =					ConfigTester(self  ,		config_class=UpperCamelCase  ,		embed_dim=37  ,		has_text_modality=UpperCamelCase					)



      def  UpperCamelCase__						(self       :       int					):

             '''simple docstring'''







             self.create_and_test_config_common_properties()
             self.config_tester.create_and_test_config_to_json_string()
             self.config_tester.create_and_test_config_to_json_file()
             self.config_tester.create_and_test_config_from_and_save_pretrained()
             self.config_tester.create_and_test_config_with_num_labels()
             self.config_tester.check_config_can_be_init_without_params()
             self.config_tester.check_config_arguments_init()



      def  UpperCamelCase__						(self       :       Union[str, Any]					):

             '''simple docstring'''







             return



      def  UpperCamelCase__						(self       :       Any					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*UpperCamelCase					)



      def  UpperCamelCase__						(self       :       Dict					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_backbone(*UpperCamelCase					)



      def  UpperCamelCase__						(self       :       List[Any]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_image_modeling(*UpperCamelCase					)



      def  UpperCamelCase__						(self       :       Dict					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*UpperCamelCase					)



      @unittest.skip(reason='''FocalNet does not use inputs_embeds'''					)
      def  UpperCamelCase__						(self       :       Dict					):

             '''simple docstring'''







             pass



      @unittest.skip(reason='''FocalNet does not use feedforward chunking'''					)
      def  UpperCamelCase__						(self       :       Dict					):

             '''simple docstring'''







             pass



      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             lowercase__     ,lowercase__      =					self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes[:-1]:
                    lowercase__      =					model_class(UpperCamelCase					)
                    self.assertIsInstance(model.get_input_embeddings()  ,		(nn.Module)					)
                    lowercase__      =					model.get_output_embeddings()
                    self.assertTrue(x is None or isinstance(UpperCamelCase  ,		nn.Linear					)					)



      def  UpperCamelCase__						(self       :       Any					):

             '''simple docstring'''







             lowercase__     ,lowercase__      =					self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes[:-1]:
                    lowercase__      =					model_class(UpperCamelCase					)
                    lowercase__      =					inspect.signature(model.forward					)
                    # signature.parameters is an OrderedDict => so arg_names order is deterministic
                    lowercase__      =					[*signature.parameters.keys()]

                    lowercase__      =					['''pixel_values''']
                    self.assertListEqual(arg_names[:1]  ,		UpperCamelCase					)



      def  UpperCamelCase__						(self       :       Optional[Any]  ,		UpperCamelCase       :       Optional[int]  ,		UpperCamelCase       :       List[Any]  ,		UpperCamelCase       :       Optional[int]  ,		UpperCamelCase       :       int					):

             '''simple docstring'''







             lowercase__      =					model_class(UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()

             with torch.no_grad():
                    lowercase__      =					model(**self._prepare_for_class(UpperCamelCase  ,		UpperCamelCase					)					)

             lowercase__      =					outputs.hidden_states

             lowercase__      =					getattr(
                 self.model_tester  ,		'''expected_num_hidden_layers'''  ,		len(self.model_tester.depths					) + 1					)
             self.assertEqual(len(UpperCamelCase					)  ,		UpperCamelCase					)

             # FocalNet has a different seq_length
             lowercase__      =					(
                 config.patch_size
                 if isinstance(config.patch_size  ,		collections.abc.Iterable					)
                 else (config.patch_size, config.patch_size)
             )

             lowercase__      =					(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

             self.assertListEqual(
                 list(hidden_states[0].shape[-2:]					)  ,		[num_patches, self.model_tester.embed_dim]  ,		)

             lowercase__      =					outputs.reshaped_hidden_states
             self.assertEqual(len(UpperCamelCase					)  ,		UpperCamelCase					)

             lowercase__     ,lowercase__     ,lowercase__     ,lowercase__      =					reshaped_hidden_states[0].shape
             lowercase__      =					(
                 reshaped_hidden_states[0].view(UpperCamelCase  ,		UpperCamelCase  ,		height * width					).permute(0  ,		2  ,		1					)
             )
             self.assertListEqual(
                 list(reshaped_hidden_states.shape[-2:]					)  ,		[num_patches, self.model_tester.embed_dim]  ,		)



      def  UpperCamelCase__						(self       :       str					):

             '''simple docstring'''







             lowercase__     ,lowercase__      =					self.model_tester.prepare_config_and_inputs_for_common()

             lowercase__      =					(
                 self.model_tester.image_size
                 if isinstance(self.model_tester.image_size  ,		collections.abc.Iterable					)
                 else (self.model_tester.image_size, self.model_tester.image_size)
             )

             for model_class in self.all_model_classes[:-1]:
                    lowercase__      =					True
                    self.check_hidden_states_output(UpperCamelCase  ,		UpperCamelCase  ,		UpperCamelCase  ,		UpperCamelCase					)

                    # check that output_hidden_states also work using config
                    del inputs_dict["output_hidden_states"]
                    lowercase__      =					True

                    self.check_hidden_states_output(UpperCamelCase  ,		UpperCamelCase  ,		UpperCamelCase  ,		UpperCamelCase					)



      def  UpperCamelCase__						(self       :       str					):

             '''simple docstring'''







             lowercase__     ,lowercase__      =					self.model_tester.prepare_config_and_inputs_for_common()
             lowercase__      =					3

             lowercase__      =					(
                 self.model_tester.image_size
                 if isinstance(self.model_tester.image_size  ,		collections.abc.Iterable					)
                 else (self.model_tester.image_size, self.model_tester.image_size)
             )
             lowercase__      =					(
                 config.patch_size
                 if isinstance(config.patch_size  ,		collections.abc.Iterable					)
                 else (config.patch_size, config.patch_size)
             )

             lowercase__      =					image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
             lowercase__      =					image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

             for model_class in self.all_model_classes[:-1]:
                    lowercase__      =					True
                    self.check_hidden_states_output(UpperCamelCase  ,		UpperCamelCase  ,		UpperCamelCase  ,		(padded_height, padded_width)					)

                    # check that output_hidden_states also work using config
                    del inputs_dict["output_hidden_states"]
                    lowercase__      =					True
                    self.check_hidden_states_output(UpperCamelCase  ,		UpperCamelCase  ,		UpperCamelCase  ,		(padded_height, padded_width)					)



      @slow
      def  UpperCamelCase__						(self       :       Tuple					):

             '''simple docstring'''







             for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    lowercase__      =					FocalNetModel.from_pretrained(UpperCamelCase					)
                    self.assertIsNotNone(UpperCamelCase					)






      def  UpperCamelCase__						(self       :       Union[str, Any]					):

             '''simple docstring'''







             lowercase__     ,lowercase__      =					self.model_tester.prepare_config_and_inputs_for_common()

             lowercase__      =					_config_zero_init(UpperCamelCase					)
             for model_class in self.all_model_classes:
                    lowercase__      =					model_class(config=UpperCamelCase					)
                    for name, param in model.named_parameters():
                           if "embeddings" not in name and param.requires_grad:
                                  self.assertIn(
                                      ((param.data.mean() * 1E9).round() / 1E9).item()  ,		[0.0, 1.0]  ,		msg=f"Parameter {name} of model {model_class} seems not properly initialized"  ,		)






@require_vision
@require_torch
class     __lowerCAmelCase		(unittest.TestCase	):





      '''simple docstring'''






      @cached_property
      def  UpperCamelCase__						(self       :       Any					):

             '''simple docstring'''







             return AutoImageProcessor.from_pretrained('''microsoft/focalnet-tiny'''					) if is_vision_available() else None






      @slow
      def  UpperCamelCase__						(self       :       Dict					):

             '''simple docstring'''







             lowercase__      =					FocalNetForImageClassification.from_pretrained('''microsoft/focalnet-tiny'''					).to(UpperCamelCase					)
             lowercase__      =					self.default_image_processor

             lowercase__      =					Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''					)
             lowercase__      =					image_processor(images=UpperCamelCase  ,		return_tensors='''pt'''					).to(UpperCamelCase					)

             # forward pass
             with torch.no_grad():
                    lowercase__      =					model(**UpperCamelCase					)

             # verify the logits
             lowercase__      =					torch.Size((1, 1000)					)
             self.assertEqual(outputs.logits.shape  ,		UpperCamelCase					)
             lowercase__      =					torch.tensor([0.21_66, -0.43_68, 0.21_91]					).to(UpperCamelCase					)
             self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,		UpperCamelCase  ,		atol=1E-4					)					)
             self.assertTrue(outputs.logits.argmax(dim=-1					).item()  ,		281					)






@require_torch
class     __lowerCAmelCase		(lowercase_    ,	unittest.TestCase	):





      '''simple docstring'''






      lowerCAmelCase__   :	List[str]       =				(FocalNetBackbone,) if is_torch_available() else ()
      lowerCAmelCase__   :	Optional[Any]       =				FocalNetConfig

      lowerCAmelCase__   :	Any       =				False



      def  UpperCamelCase__						(self       :       List[str]					):

             '''simple docstring'''







             lowercase__      =					FocalNetModelTester(self					)








'''simple docstring'''

import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch

def _SCREAMING_SNAKE_CASE (A			)			->			Dict:






       """simple docstring"""







       lowercase__      =					os.path.join(args.tf_model_dir    ,					'''parameters.json'''			)
       lowercase__      =					json.loads(open(A			).read()			)
       if not params:
              raise ValueError(
                  f"It seems that the json file at {parameter_file} is empty. Make sure you have a correct json file."			)
       if not args.output.endswith('''.pt'''			):
              lowercase__      =					args.output + '''.pt'''
       lowercase__      =					OrderedDict()
       with tf.device('''/CPU:0'''			):
              lowercase__      =					tf.train.load_checkpoint(args.tf_model_dir			)
              lowercase__      =					reader.get_variable_to_shape_map()
              for key_name in shapes.keys():
                     lowercase__      =					reader.get_tensor(A			).astype(np.floataa			)
                     if key_name.endswith('''/adam_m'''			) or key_name.endswith('''/adam_v'''			):
                            continue
                     if key_name.startswith('''pasts/'''			):
                            if key_name.startswith('''pasts/mlp'''			):
                                   lowercase__      =					int(key_name[9]			)
                            elif key_name.startswith('''pasts/out'''			):
                                   lowercase__      =					8
                            lowercase__      =					'''model.sqout.%d.weight''' % (player * 2)  # enter to nn.Sequencial with Tanh, so 2 at a time
                            lowercase__      =					vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
                            lowercase__      =					torch.tensor(A			)
                     elif key_name.startswith('''model/moe'''			):
                            lowercase__      =					int(key_name[9:].split('''/'''			)[0]			)
                            if key_name.endswith('''/switch_gating/kernel'''			):
                                   lowercase__      =					'''model.blocks.%d.feed_forward.mlp.router.classifier.weight''' % player
                                   lowercase__      =					vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
                                   lowercase__      =					torch.tensor(A			)
                            elif key_name.endswith('''/softmlp/kernel'''			):
                                   lowercase__      =					'''model.blocks.%d.feed_forward.soft_bypass_mlp.weight''' % player
                                   lowercase__      =					vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
                                   lowercase__      =					torch.tensor(A			)
                            elif key_name.endswith('''/wo/kernel'''			) or key_name.endswith('''/wi/kernel'''			):
                                   lowercase__      =					key_name[-9:-7]
                                   for i in range(16			):
                                          lowercase__      =					'''model.blocks.%d.feed_forward.mlp.experts.expert_%d.%s.weight''' % (player, i, nlayer)
                                          lowercase__      =					(
                                              vnp[i].transpose([1, 0]			).copy()
                                          )  # In Mesh-Tensorflow, it is one array, so it is divided
                                          lowercase__      =					torch.tensor(A			)
                     elif key_name.startswith('''model/mlp'''			):
                            lowercase__      =					int(key_name[9:].split('''/'''			)[0]			)
                            if key_name.endswith('''/p1/kernel'''			):
                                   lowercase__      =					'''model.blocks.%d.feed_forward.mlp.wi.weight''' % player
                                   lowercase__      =					vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
                                   lowercase__      =					torch.tensor(A			)
                            elif key_name.endswith('''/p1/bias'''			):
                                   lowercase__      =					'''model.blocks.%d.feed_forward.mlp.wi.bias''' % player
                                   lowercase__      =					vnp.copy()  # same because it is one dimensional
                                   lowercase__      =					torch.tensor(A			)
                            elif key_name.endswith('''/p2/kernel'''			):
                                   lowercase__      =					'''model.blocks.%d.feed_forward.mlp.wo.weight''' % player
                                   lowercase__      =					vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
                                   lowercase__      =					torch.tensor(A			)
                            elif key_name.endswith('''/p2/bias'''			):
                                   lowercase__      =					'''model.blocks.%d.feed_forward.mlp.wo.bias''' % player
                                   lowercase__      =					vnp.copy()  # same because it is one dimensional
                                   lowercase__      =					torch.tensor(A			)
                     elif key_name.startswith('''model/ln'''			):
                            lowercase__      =					int(key_name[8:].split('''/'''			)[0]			)
                            if key_name.endswith('''/b'''			):
                                   lowercase__      =					'''model.blocks.%d.feed_forward.norm.bias''' % player
                                   lowercase__      =					vnp.copy()  # same because it is one dimensional
                                   lowercase__      =					torch.tensor(A			)
                            elif key_name.endswith('''/g'''			):
                                   lowercase__      =					'''model.blocks.%d.feed_forward.norm.weight''' % player
                                   lowercase__      =					vnp.copy()  # same because it is one dimensional
                                   lowercase__      =					torch.tensor(A			)
                     elif key_name.startswith('''model/att'''			):
                            lowercase__      =					int(key_name[9:].split('''/'''			)[0]			)
                            if key_name.endswith('''/qkv/kernel'''			):
                                   lowercase__      =					vnp.copy()  # Compute same dimension as Mesh-tensorflow using einsum
                                   lowercase__      =					state[:, 0, :, :]
                                   lowercase__      =					state[:, 1, :, :]
                                   lowercase__      =					state[:, 2, :, :]
                                   lowercase__      =					(
                                       state_q.reshape([state_q.shape[0], state_q.shape[1] * state_q.shape[2]]			)
                                       .transpose([1, 0]			)
                                       .copy()
                                   )  # Mesh-Tensorflow is a diagonal matrix
                                   lowercase__      =					(
                                       state_k.reshape([state_k.shape[0], state_k.shape[1] * state_k.shape[2]]			)
                                       .transpose([1, 0]			)
                                       .copy()
                                   )  # Mesh-Tensorflow is a diagonal matrix
                                   lowercase__      =					(
                                       state_v.reshape([state_v.shape[0], state_v.shape[1] * state_v.shape[2]]			)
                                       .transpose([1, 0]			)
                                       .copy()
                                   )  # Mesh-Tensorflow is a diagonal matrix
                                   lowercase__      =					'''model.blocks.%d.self_attn.self_attn.q_proj.weight''' % player
                                   lowercase__      =					torch.tensor(A			)
                                   lowercase__      =					'''model.blocks.%d.self_attn.self_attn.k_proj.weight''' % player
                                   lowercase__      =					torch.tensor(A			)
                                   lowercase__      =					'''model.blocks.%d.self_attn.self_attn.v_proj.weight''' % player
                                   lowercase__      =					torch.tensor(A			)
                            elif key_name.endswith('''/o/kernel'''			):
                                   lowercase__      =					'''model.blocks.%d.self_attn.self_attn.out_proj.weight''' % player
                                   lowercase__      =					(
                                       vnp.reshape([vnp.shape[0] * vnp.shape[1], vnp.shape[2]]			).transpose([1, 0]			).copy()
                                   )  # Mesh-Tensorflow is a diagonal matrix
                                   lowercase__      =					torch.tensor(A			)
                     elif key_name.startswith('''model/an'''			):
                            lowercase__      =					int(key_name[8:].split('''/'''			)[0]			)
                            if key_name.endswith('''/b'''			):
                                   lowercase__      =					'''model.blocks.%d.self_attn.norm.bias''' % player
                                   lowercase__      =					vnp.copy()  # same because it is one dimensional
                                   lowercase__      =					torch.tensor(A			)
                            elif key_name.endswith('''/g'''			):
                                   lowercase__      =					'''model.blocks.%d.self_attn.norm.weight''' % player
                                   lowercase__      =					vnp.copy()  # same because it is one dimensional
                                   lowercase__      =					torch.tensor(A			)
                     elif (
                         key_name.startswith('''model/wte'''			)
                         or key_name.startswith('''model/wpe'''			)
                         or key_name.startswith('''model/ete'''			)
                     ):
                            lowercase__      =					{'''wte''': '''embed_tokens''', '''wpe''': '''position_embeddings''', '''ete''': '''extra_position_embeddings'''}[
                                key_name[-3:]
                            ]
                            lowercase__      =					'''model.%s.weight''' % nlayer
                            lowercase__      =					vnp.copy()  # same in embedded
                            lowercase__      =					torch.tensor(A			)
                            if key_name.startswith('''model/wte'''			):
                                   lowercase__      =					'''lm_head.weight'''
                                   lowercase__      =					vnp.copy()  # same in embedded
                                   lowercase__      =					torch.tensor(A			)
                     elif key_name.startswith('''model/wob'''			):
                            lowercase__      =					'''final_logits_bias'''
                            lowercase__      =					vnp.copy()  # same in embedded
                            lowercase__      =					state.reshape((1, -1)			)
                            lowercase__      =					torch.tensor(A			)
                     elif key_name == "model/dense/kernel":
                            lowercase__      =					'''model.last_project.weight'''
                            lowercase__      =					vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
                            lowercase__      =					torch.tensor(A			)
                     elif key_name == "model/dense_1/bias":
                            lowercase__      =					'''model.last_project.bias'''
                            lowercase__      =					vnp.copy()  # same because it is one dimensional
                            lowercase__      =					torch.tensor(A			)
       torch.save(A    ,					args.output			)


if __name__ == "__main__":
 lowerCamelCase				:  Optional[int]    		=      argparse.ArgumentParser(
     description='model converter.', formatter_class=argparse.ArgumentDefaultsHelpFormatter
 )
 parser.add_argument('--tf_model_dir', metavar='PATH', type=str, required=True, help='import model')
 parser.add_argument('--output', metavar='PATH', type=str, required=True, help='output model')
 lowerCamelCase				:  Dict    		=      parser.parse_args()
 convert_tf_gptsan_to_pt(args)




