'''simple docstring'''

import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
  import torch
  from torch import nn

  from transformers import ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
  from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import ViTImageProcessor



class       UpperCAmelCase  :





       '''simple docstring'''



       def __init__(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase=13     ,       __lowerCAmelCase=30     ,       __lowerCAmelCase=2     ,       __lowerCAmelCase=3     ,       __lowerCAmelCase=True     ,       __lowerCAmelCase=True     ,       __lowerCAmelCase=32     ,       __lowerCAmelCase=5     ,       __lowerCAmelCase=4     ,       __lowerCAmelCase=37     ,       __lowerCAmelCase="gelu"     ,       __lowerCAmelCase=0.1     ,       __lowerCAmelCase=0.1     ,       __lowerCAmelCase=10     ,       __lowerCAmelCase=0.0_2     ,       __lowerCAmelCase=None     ,       __lowerCAmelCase=2     ,       )  ->  Dict:
        lowercase__    :   List[Any]						= parent
        lowercase__    :   List[Any]						= batch_size
        lowercase__    :   List[Any]						= image_size
        lowercase__    :   str						= patch_size
        lowercase__    :   Tuple						= num_channels
        lowercase__    :   Tuple						= is_training
        lowercase__    :   Tuple						= use_labels
        lowercase__    :   List[str]						= hidden_size
        lowercase__    :   Union[str, Any]						= num_hidden_layers
        lowercase__    :   Optional[Any]						= num_attention_heads
        lowercase__    :   Dict						= intermediate_size
        lowercase__    :   Tuple						= hidden_act
        lowercase__    :   List[str]						= hidden_dropout_prob
        lowercase__    :   int						= attention_probs_dropout_prob
        lowercase__    :   List[Any]						= type_sequence_label_size
        lowercase__    :   Dict						= initializer_range
        lowercase__    :   Dict						= scope
        lowercase__    :   Optional[Any]						= encoder_stride

        # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
        lowercase__    :   Union[str, Any]						= (image_size // patch_size) ** 2
        lowercase__    :   Any						= num_patches + 1

       def _lowerCAmelCase(						self			)  ->  List[Any]:
        lowercase__    :   Dict						= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

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

        lowercase__    :   int						= self.get_config()

        return config, pixel_values, labels

       def _lowerCAmelCase(						self			)  ->  Optional[Any]:
        return ViTConfig(
            image_size=self.image_size     ,       patch_size=self.patch_size     ,       num_channels=self.num_channels     ,       hidden_size=self.hidden_size     ,       num_hidden_layers=self.num_hidden_layers     ,       num_attention_heads=self.num_attention_heads     ,       intermediate_size=self.intermediate_size     ,       hidden_act=self.hidden_act     ,       hidden_dropout_prob=self.hidden_dropout_prob     ,       attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,       is_decoder=__lowerCAmelCase     ,       initializer_range=self.initializer_range     ,       encoder_stride=self.encoder_stride     ,       )

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  Tuple:
        lowercase__    :   List[Any]						= ViTModel(config=__lowerCAmelCase			)
        model.to(__lowerCAmelCase			)
        model.eval()
        lowercase__    :   Any						= model(__lowerCAmelCase			)
        self.parent.assertEqual(result.last_hidden_state.shape     ,       (self.batch_size, self.seq_length, self.hidden_size)			)

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  List[Any]:
        lowercase__    :   Tuple						= ViTForMaskedImageModeling(config=__lowerCAmelCase			)
        model.to(__lowerCAmelCase			)
        model.eval()
        lowercase__    :   str						= model(__lowerCAmelCase			)
        self.parent.assertEqual(
            result.reconstruction.shape     ,       (self.batch_size, self.num_channels, self.image_size, self.image_size)			)

        # test greyscale images
        lowercase__    :   Any						= 1
        lowercase__    :   Optional[Any]						= ViTForMaskedImageModeling(__lowerCAmelCase			)
        model.to(__lowerCAmelCase			)
        model.eval()

        lowercase__    :   str						= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
        lowercase__    :   int						= model(__lowerCAmelCase			)
        self.parent.assertEqual(result.reconstruction.shape     ,       (self.batch_size, 1, self.image_size, self.image_size)			)

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  List[Any]:
        lowercase__    :   Any						= self.type_sequence_label_size
        lowercase__    :   Dict						= ViTForImageClassification(__lowerCAmelCase			)
        model.to(__lowerCAmelCase			)
        model.eval()
        lowercase__    :   int						= model(__lowerCAmelCase     ,       labels=__lowerCAmelCase			)
        self.parent.assertEqual(result.logits.shape     ,       (self.batch_size, self.type_sequence_label_size)			)

        # test greyscale images
        lowercase__    :   List[str]						= 1
        lowercase__    :   Optional[int]						= ViTForImageClassification(__lowerCAmelCase			)
        model.to(__lowerCAmelCase			)
        model.eval()

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


       def _lowerCAmelCase(						self			)  ->  Dict:
        lowercase__    :   Union[str, Any]						= self.prepare_config_and_inputs()
        (
            (
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							
        )    :   Union[str, Any]						= config_and_inputs
        lowercase__    :   List[str]						= {'''pixel_values''': pixel_values}
        return config, inputs_dict



@require_torch
class       UpperCAmelCase  (						a__   ,						a__   ,						unittest.TestCase ):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE 		=   (
           (
               ViTModel,
               ViTForImageClassification,
               ViTForMaskedImageModeling,
           )
           if is_torch_available()
           else ()
       )
       SCREAMING_SNAKE_CASE 		=   (
           {"feature-extraction": ViTModel, "image-classification": ViTForImageClassification}
           if is_torch_available()
           else {}
       )
       SCREAMING_SNAKE_CASE 		=   True

       SCREAMING_SNAKE_CASE 		=   False
       SCREAMING_SNAKE_CASE 		=   False
       SCREAMING_SNAKE_CASE 		=   False

       def _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        lowercase__    :   Optional[int]						= ViTModelTester(self			)
        lowercase__    :   int						= ConfigTester(self     ,       config_class=__lowerCAmelCase     ,       has_text_modality=__lowerCAmelCase     ,       hidden_size=37			)

       def _lowerCAmelCase(						self			)  ->  List[Any]:
        self.config_tester.run_common_tests()

       @unittest.skip(reason='''ViT does not use inputs_embeds'''			)
       def _lowerCAmelCase(						self			)  ->  Dict:
        pass

       def _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        lowercase__			,							lowercase__    :   List[Any]						= self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
         lowercase__    :   Any						= model_class(__lowerCAmelCase			)
         self.assertIsInstance(model.get_input_embeddings()     ,       (nn.Module)			)
         lowercase__    :   Union[str, Any]						= model.get_output_embeddings()
         self.assertTrue(x is None or isinstance(__lowerCAmelCase     ,       nn.Linear			)			)

       def _lowerCAmelCase(						self			)  ->  Optional[int]:
        lowercase__			,							lowercase__    :   int						= self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
         lowercase__    :   List[Any]						= model_class(__lowerCAmelCase			)
         lowercase__    :   List[str]						= inspect.signature(model.forward			)
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         lowercase__    :   Optional[int]						= [*signature.parameters.keys()]

         lowercase__    :   Optional[Any]						= ['''pixel_values''']
         self.assertListEqual(arg_names[:1]     ,       __lowerCAmelCase			)

       def _lowerCAmelCase(						self			)  ->  str:
        lowercase__    :   Optional[Any]						= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCAmelCase			)

       def _lowerCAmelCase(						self			)  ->  Dict:
        lowercase__    :   Dict						= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_image_modeling(*__lowerCAmelCase			)

       def _lowerCAmelCase(						self			)  ->  int:
        lowercase__    :   List[Any]						= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*__lowerCAmelCase			)


       @slow
       def _lowerCAmelCase(						self			)  ->  List[Any]:
        for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowercase__    :   Optional[Any]						= ViTModel.from_pretrained(__lowerCAmelCase			)
         self.assertIsNotNone(__lowerCAmelCase			)


def       __UpperCamelCase	(       ):
 lowercase__    :   Dict						= Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''					)
 return image





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





       '''simple docstring'''



       @cached_property
       def _lowerCAmelCase(						self			)  ->  List[str]:
        return ViTImageProcessor.from_pretrained('''google/vit-base-patch16-224'''			) if is_vision_available() else None

       @slow
       def _lowerCAmelCase(						self			)  ->  List[Any]:
        lowercase__    :   Union[str, Any]						= ViTForImageClassification.from_pretrained('''google/vit-base-patch16-224'''			).to(__lowerCAmelCase			)

        lowercase__    :   str						= self.default_image_processor
        lowercase__    :   str						= prepare_img()
        lowercase__    :   Tuple						= image_processor(images=__lowerCAmelCase     ,       return_tensors='''pt'''			).to(__lowerCAmelCase			)

        # forward pass
        with torch.no_grad():
         lowercase__    :   List[str]						= model(**__lowerCAmelCase			)

        # verify the logits
        lowercase__    :   List[Any]						= torch.Size((1, 1000)			)
        self.assertEqual(outputs.logits.shape     ,       __lowerCAmelCase			)

        lowercase__    :   Optional[Any]						= torch.tensor([-0.2_7_4_4, 0.8_2_1_5, -0.0_8_3_6]			).to(__lowerCAmelCase			)

        self.assertTrue(torch.allclose(outputs.logits[0, :3]     ,       __lowerCAmelCase     ,       atol=1E-4			)			)

       @slow
       def _lowerCAmelCase(						self			)  ->  Tuple:
        # ViT models have an `interpolate_pos_encoding` argument in their forward method,
        # allowing to interpolate the pre-trained position embeddings in order to use
        # the model on higher resolutions. The DINO model by Facebook AI leverages this
        # to visualize self-attention on higher resolution images.
        lowercase__    :   int						= ViTModel.from_pretrained('''facebook/dino-vits8'''			).to(__lowerCAmelCase			)

        lowercase__    :   Optional[int]						= ViTImageProcessor.from_pretrained('''facebook/dino-vits8'''     ,       size=480			)
        lowercase__    :   int						= prepare_img()
        lowercase__    :   int						= image_processor(images=__lowerCAmelCase     ,       return_tensors='''pt'''			)
        lowercase__    :   Optional[int]						= inputs.pixel_values.to(__lowerCAmelCase			)

        # forward pass
        with torch.no_grad():
         lowercase__    :   Union[str, Any]						= model(__lowerCAmelCase     ,       interpolate_pos_encoding=__lowerCAmelCase			)

        # verify the logits
        lowercase__    :   str						= torch.Size((1, 3601, 384)			)
        self.assertEqual(outputs.last_hidden_state.shape     ,       __lowerCAmelCase			)

        lowercase__    :   List[str]						= torch.tensor(
            [[4.2_3_4_0, 4.3_9_0_6, -6.6_6_9_2], [4.5_4_6_3, 1.8_9_2_8, -6.7_2_5_7], [4.4_4_2_9, 0.8_4_9_6, -5.8_5_8_5]]			).to(__lowerCAmelCase			)

        self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]     ,       __lowerCAmelCase     ,       atol=1E-4			)			)


       @slow
       @require_accelerate
       @require_torch_gpu
       def _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        lowercase__    :   Tuple						= ViTModel.from_pretrained('''facebook/dino-vits8'''     ,       torch_dtype=torch.floataa     ,       device_map='''auto'''			)
        lowercase__    :   str						= self.default_image_processor

        lowercase__    :   str						= prepare_img()
        lowercase__    :   Optional[Any]						= image_processor(images=__lowerCAmelCase     ,       return_tensors='''pt'''			)
        lowercase__    :   Optional[Any]						= inputs.pixel_values.to(__lowerCAmelCase			)

        # forward pass to make sure inference works in fp16
        with torch.no_grad():
         lowercase__    :   Optional[Any]						= model(__lowerCAmelCase			)


'''simple docstring'''

import argparse
import collections

import torch
from flax import traverse_util
from tax import checkpoints

from transformers import TaConfig, TaEncoderModel, TaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase="attention"					):
 lowercase__    :   Dict						= params[F"""{prefix}/layers_{i}/{layer_name}/key/kernel"""]
 lowercase__    :   Dict						= params[F"""{prefix}/layers_{i}/{layer_name}/out/kernel"""]
 lowercase__    :   int						= params[F"""{prefix}/layers_{i}/{layer_name}/query/kernel"""]
 lowercase__    :   Any						= params[F"""{prefix}/layers_{i}/{layer_name}/value/kernel"""]
 return k, o, q, v
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase=False					):
 if split_mlp_wi:
  lowercase__    :   Tuple						= params[F"""{prefix}/layers_{i}/mlp/wi_0/kernel"""]
  lowercase__    :   Tuple						= params[F"""{prefix}/layers_{i}/mlp/wi_1/kernel"""]
  lowercase__    :   Dict						= (wi_a, wi_a)
 else:
  lowercase__    :   List[Any]						= params[F"""{prefix}/layers_{i}/mlp/wi/kernel"""]

 lowercase__    :   int						= params[F"""{prefix}/layers_{i}/mlp/wo/kernel"""]
 return wi, wo
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase					):
 return params[F"""{prefix}/layers_{i}/{layer_name}/scale"""]
def       __UpperCamelCase	(       UpperCAmelCase		,					*, UpperCAmelCase		,					UpperCAmelCase					):
 lowercase__    :   Union[str, Any]						= traverse_util.flatten_dict(variables['''target''']					)
 lowercase__    :   Any						= {'''/'''.join(UpperCAmelCase					): v for k, v in old.items()}

 # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
 lowercase__    :   List[str]						= '''encoder/layers_0/mlp/wi_0/kernel''' in old
 print('''Split MLP:'''		,					UpperCAmelCase					)

 lowercase__    :   Union[str, Any]						= collections.OrderedDict()

 # Shared embeddings.
 lowercase__    :   str						= old['''token_embedder/embedding''']

 # Encoder.
 for i in range(UpperCAmelCase					):
  # Block i, layer 0 (Self Attention).
  lowercase__    :   List[Any]						= tax_layer_norm_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''encoder'''		,					'''pre_attention_layer_norm'''					)
  lowercase__			,							lowercase__			,							lowercase__			,							lowercase__    :   Tuple						= tax_attention_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''encoder'''		,					'''attention'''					)
  lowercase__    :   Dict						= layer_norm
  lowercase__    :   Any						= k.T
  lowercase__    :   Optional[int]						= o.T
  lowercase__    :   Optional[int]						= q.T
  lowercase__    :   Any						= v.T

  # Block i, layer 1 (MLP).
  lowercase__    :   List[str]						= tax_layer_norm_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''encoder'''		,					'''pre_mlp_layer_norm'''					)
  lowercase__			,							lowercase__    :   List[Any]						= tax_mlp_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''encoder'''		,					UpperCAmelCase					)
  lowercase__    :   int						= layer_norm
  if split_mlp_wi:
   lowercase__    :   Union[str, Any]						= wi[0].T
   lowercase__    :   List[str]						= wi[1].T
  else:
   lowercase__    :   int						= wi.T
  lowercase__    :   List[str]						= wo.T

 lowercase__    :   Dict						= old[
     '''encoder/relpos_bias/rel_embedding'''
 ].T
 lowercase__    :   str						= old['''encoder/encoder_norm/scale''']

 if not is_encoder_only:
  # Decoder.
  for i in range(UpperCAmelCase					):
   # Block i, layer 0 (Self Attention).
   lowercase__    :   Optional[Any]						= tax_layer_norm_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					'''pre_self_attention_layer_norm'''					)
   lowercase__			,							lowercase__			,							lowercase__			,							lowercase__    :   Optional[Any]						= tax_attention_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					'''self_attention'''					)
   lowercase__    :   List[str]						= layer_norm
   lowercase__    :   List[Any]						= k.T
   lowercase__    :   List[Any]						= o.T
   lowercase__    :   List[str]						= q.T
   lowercase__    :   Optional[Any]						= v.T

   # Block i, layer 1 (Cross Attention).
   lowercase__    :   List[Any]						= tax_layer_norm_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					'''pre_cross_attention_layer_norm'''					)
   lowercase__			,							lowercase__			,							lowercase__			,							lowercase__    :   Any						= tax_attention_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					'''encoder_decoder_attention'''					)
   lowercase__    :   Any						= layer_norm
   lowercase__    :   Optional[int]						= k.T
   lowercase__    :   Optional[Any]						= o.T
   lowercase__    :   Tuple						= q.T
   lowercase__    :   List[Any]						= v.T

   # Block i, layer 2 (MLP).
   lowercase__    :   List[str]						= tax_layer_norm_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					'''pre_mlp_layer_norm'''					)
   lowercase__			,							lowercase__    :   List[Any]						= tax_mlp_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					UpperCAmelCase					)
   lowercase__    :   Optional[int]						= layer_norm
   if split_mlp_wi:
    lowercase__    :   List[str]						= wi[0].T
    lowercase__    :   str						= wi[1].T
   else:
    lowercase__    :   str						= wi.T
   lowercase__    :   Optional[Any]						= wo.T

  lowercase__    :   Any						= old['''decoder/decoder_norm/scale''']
  lowercase__    :   Optional[Any]						= old[
      '''decoder/relpos_bias/rel_embedding'''
  ].T

  # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
  if "decoder/logits_dense/kernel" in old:
   lowercase__    :   Dict						= old['''decoder/logits_dense/kernel'''].T

 return new
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase					):
 lowercase__    :   int						= collections.OrderedDict([(k, torch.from_numpy(v.copy()					)) for (k, v) in converted_params.items()]					)

 # Add what is missing.
 if "encoder.embed_tokens.weight" not in state_dict:
  lowercase__    :   Optional[int]						= state_dict['''shared.weight''']

 if not is_encoder_only:
  if "decoder.embed_tokens.weight" not in state_dict:
   lowercase__    :   List[Any]						= state_dict['''shared.weight''']

  if "lm_head.weight" not in state_dict:  # For old 1.0 models.
   print('''Using shared word embeddings as lm_head.'''					)
   lowercase__    :   Union[str, Any]						= state_dict['''shared.weight''']

 return state_dict
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase					):
 lowercase__    :   List[str]						= checkpoints.load_tax_checkpoint(UpperCAmelCase					)
 lowercase__    :   str						= convert_tax_to_pytorch(UpperCAmelCase		,					num_layers=config.num_layers		,					is_encoder_only=UpperCAmelCase					)
 lowercase__    :   List[Any]						= make_state_dict(UpperCAmelCase		,					UpperCAmelCase					)
 model.load_state_dict(UpperCAmelCase		,					strict=UpperCAmelCase					)
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase = False					):
 lowercase__    :   Dict						= TaConfig.from_json_file(UpperCAmelCase					)
 print(F"""Building PyTorch model from configuration: {config}"""					)
 # Non-v1.1 checkpoints could also use T5Model, but this works for all.
 # The v1.0 checkpoints will simply have an LM head that is the word embeddings.
 if is_encoder_only:
  lowercase__    :   List[str]						= TaEncoderModel(UpperCAmelCase					)
 else:
  lowercase__    :   Tuple						= TaForConditionalGeneration(UpperCAmelCase					)

 # Load weights from tf checkpoint
 load_tax_weights_in_ta(UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase					)

 # Save pytorch-model
 print(F"""Save PyTorch model to {pytorch_dump_path}"""					)
 model.save_pretrained(UpperCAmelCase					)

 # Verify that we can load the checkpoint.
 model.from_pretrained(UpperCAmelCase					)
 print('''Done'''					)


if __name__ == "__main__":
  __a:			Optional[Any]	       =       argparse.ArgumentParser(description="""Converts a native T5X checkpoint into a PyTorch checkpoint.""")
  # Required parameters
  parser.add_argument(
      """--t5x_checkpoint_path""", default=None, type=str, required=True, help="""Path to the T5X checkpoint."""
  )
  parser.add_argument(
      """--config_file""",
      default=None,
      type=str,
      required=True,
      help="""The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.""",
  )
  parser.add_argument(
      """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
  )
  parser.add_argument(
      """--is_encoder_only""", action="""store_true""", help="""Check if the model is encoder-decoder model""", default=False
  )
  __a:			Tuple	       =       parser.parse_args()
  convert_tax_checkpoint_to_pytorch(
      args.tax_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
  )

