




from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
  from transformers import (
      TFAutoModel,
      TFAutoModelForCausalLM,
      TFAutoModelForMaskedLM,
      TFAutoModelForPreTraining,
      TFAutoModelForQuestionAnswering,
      TFAutoModelForSeqaSeqLM,
      TFAutoModelForSequenceClassification,
      TFAutoModelForTableQuestionAnswering,
      TFAutoModelForTokenClassification,
      TFAutoModelWithLMHead,
      TFBertForMaskedLM,
      TFBertForPreTraining,
      TFBertForQuestionAnswering,
      TFBertForSequenceClassification,
      TFBertModel,
      TFFunnelBaseModel,
      TFFunnelModel,
      TFGPTaLMHeadModel,
      TFRobertaForMaskedLM,
      TFTaForConditionalGeneration,
      TFTapasForQuestionAnswering,
  )
  from transformers.models.auto.modeling_tf_auto import (
      TF_MODEL_FOR_CAUSAL_LM_MAPPING,
      TF_MODEL_FOR_MASKED_LM_MAPPING,
      TF_MODEL_FOR_PRETRAINING_MAPPING,
      TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
      TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
      TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
      TF_MODEL_MAPPING,
  )
  from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
  from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
  from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
  from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST


class     UpperCamelCase_   (			UpperCAmelCase__					):

    '''simple docstring'''






    UpperCAmelCase__							  =				'''new-model'''


if is_tf_available():


  class     UpperCamelCase_   (			UpperCAmelCase__					):

      '''simple docstring'''






      UpperCAmelCase__							  =				NewModelConfig


@require_tf
class     UpperCamelCase_   (			unittest.TestCase					):

    '''simple docstring'''






    @slow
    def 	SCREAMING_SNAKE_CASE							(      self				:     str)						->Union[str, Any]:





           '''simple docstring'''

           A__							=   '''bert-base-cased'''
           A__							=   AutoConfig.from_pretrained(UpperCAmelCase__)
           self.assertIsNotNone(UpperCAmelCase__)
           self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

           A__							=   TFAutoModel.from_pretrained(UpperCAmelCase__)
           self.assertIsNotNone(UpperCAmelCase__)
           self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)


    @slow
    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[Any])						->List[Any]:





           '''simple docstring'''

           A__							=   '''bert-base-cased'''
           A__							=   AutoConfig.from_pretrained(UpperCAmelCase__)
           self.assertIsNotNone(UpperCAmelCase__)
           self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

           A__							=   TFAutoModelForPreTraining.from_pretrained(UpperCAmelCase__)
           self.assertIsNotNone(UpperCAmelCase__)
           self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)


    @slow
    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[Any])						->Dict:





           '''simple docstring'''

           for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A__							=   AutoConfig.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

                  A__							=   TFAutoModelForCausalLM.from_pretrained(UpperCAmelCase__)
                  A__			,							A__							=   TFAutoModelForCausalLM.from_pretrained(UpperCAmelCase__ ,			output_loading_info=UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)


    @slow
    def 	SCREAMING_SNAKE_CASE							(      self				:     List[Any])						->str:





           '''simple docstring'''

           for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A__							=   AutoConfig.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

                  A__							=   TFAutoModelWithLMHead.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)


    @slow
    def 	SCREAMING_SNAKE_CASE							(      self				:     Dict)						->List[str]:





           '''simple docstring'''

           for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A__							=   AutoConfig.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

                  A__							=   TFAutoModelForMaskedLM.from_pretrained(UpperCAmelCase__)
                  A__			,							A__							=   TFAutoModelForMaskedLM.from_pretrained(UpperCAmelCase__ ,			output_loading_info=UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)


    @slow
    def 	SCREAMING_SNAKE_CASE							(      self				:     List[Any])						->Tuple:





           '''simple docstring'''

           for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A__							=   AutoConfig.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

                  A__							=   TFAutoModelForSeqaSeqLM.from_pretrained(UpperCAmelCase__)
                  A__			,							A__							=   TFAutoModelForSeqaSeqLM.from_pretrained(UpperCAmelCase__ ,			output_loading_info=UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)


    @slow
    def 	SCREAMING_SNAKE_CASE							(      self				:     Any)						->List[str]:





           '''simple docstring'''

           for model_name in ["bert-base-uncased"]:
                  A__							=   AutoConfig.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

                  A__							=   TFAutoModelForSequenceClassification.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)


    @slow
    def 	SCREAMING_SNAKE_CASE							(      self				:     str)						->Optional[Any]:





           '''simple docstring'''

           for model_name in ["bert-base-uncased"]:
                  A__							=   AutoConfig.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

                  A__							=   TFAutoModelForQuestionAnswering.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)


    @slow
    @require_tensorflow_probability
    def 	SCREAMING_SNAKE_CASE							(      self				:     List[Any])						->int:





           '''simple docstring'''

           for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
                  A__							=   AutoConfig.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

                  A__							=   TFAutoModelForTableQuestionAnswering.from_pretrained(UpperCAmelCase__)
                  A__			,							A__							=   TFAutoModelForTableQuestionAnswering.from_pretrained(
                      UpperCAmelCase__ ,			output_loading_info=UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)


    def 	SCREAMING_SNAKE_CASE							(      self				:     Dict)						->str:





           '''simple docstring'''

           A__							=   TFAutoModelWithLMHead.from_pretrained(UpperCAmelCase__)
           self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)
           self.assertEqual(model.num_parameters() ,			14_410)
           self.assertEqual(model.num_parameters(only_trainable=UpperCAmelCase__) ,			14_410)


    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[int])						->Any:





           '''simple docstring'''

           A__							=   TFAutoModelWithLMHead.from_pretrained(UpperCAmelCase__)
           self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)
           self.assertEqual(model.num_parameters() ,			14_410)
           self.assertEqual(model.num_parameters(only_trainable=UpperCAmelCase__) ,			14_410)


    def 	SCREAMING_SNAKE_CASE							(      self				:     Any)						->Dict:





           '''simple docstring'''

           A__							=   TFAutoModel.from_pretrained('''sgugger/funnel-random-tiny''')
           self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

           A__							=   copy.deepcopy(model.config)
           A__							=   ['''FunnelBaseModel''']
           A__							=   TFAutoModel.from_config(UpperCAmelCase__)
           self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

           with tempfile.TemporaryDirectory() as tmp_dir:
                  model.save_pretrained(UpperCAmelCase__)
                  A__							=   TFAutoModel.from_pretrained(UpperCAmelCase__)
                  self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)


    def 	SCREAMING_SNAKE_CASE							(      self				:     str)						->Dict:





           '''simple docstring'''

           try:
                  AutoConfig.register('''new-model''' ,			UpperCAmelCase__)

                  A__							=   [
                      TFAutoModel,
                      TFAutoModelForCausalLM,
                      TFAutoModelForMaskedLM,
                      TFAutoModelForPreTraining,
                      TFAutoModelForQuestionAnswering,
                      TFAutoModelForSequenceClassification,
                      TFAutoModelForTokenClassification,
                  ]

                  for auto_class in auto_classes:
                         with self.subTest(auto_class.__name__):
                                # Wrong config class will raise an error
                                with self.assertRaises(UpperCAmelCase__):
                                       auto_class.register(UpperCAmelCase__ ,			UpperCAmelCase__)
                                auto_class.register(UpperCAmelCase__ ,			UpperCAmelCase__)
                                # Trying to register something existing in the Transformers library will raise an error
                                with self.assertRaises(UpperCAmelCase__):
                                       auto_class.register(UpperCAmelCase__ ,			UpperCAmelCase__)

                                # Now that the config is registered, it can be used as any other config with the auto-API
                                A__							=   BertModelTester(self).get_config()
                                A__							=   NewModelConfig(**tiny_config.to_dict())
                                A__							=   auto_class.from_config(UpperCAmelCase__)
                                self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

                                with tempfile.TemporaryDirectory() as tmp_dir:
                                       model.save_pretrained(UpperCAmelCase__)
                                       A__							=   auto_class.from_pretrained(UpperCAmelCase__)
                                       self.assertIsInstance(UpperCAmelCase__ ,			UpperCAmelCase__)

           finally:
                  if "new-model" in CONFIG_MAPPING._extra_content:
                         del CONFIG_MAPPING._extra_content["new-model"]
                  for mapping in (
                      TF_MODEL_MAPPING,
                      TF_MODEL_FOR_PRETRAINING_MAPPING,
                      TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
                      TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
                      TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
                      TF_MODEL_FOR_CAUSAL_LM_MAPPING,
                      TF_MODEL_FOR_MASKED_LM_MAPPING,
                  ):
                         if NewModelConfig in mapping._extra_content:
                                del mapping._extra_content[NewModelConfig]


    def 	SCREAMING_SNAKE_CASE							(      self				:     List[str])						->int:





           '''simple docstring'''

           with self.assertRaisesRegex(
               UpperCAmelCase__ ,			'''bert-base is not a local folder and is not a valid model identifier'''):
                  A__							=   TFAutoModel.from_pretrained('''bert-base''')


    def 	SCREAMING_SNAKE_CASE							(      self				:     str)						->Tuple:





           '''simple docstring'''

           with self.assertRaisesRegex(
               UpperCAmelCase__ ,			R'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''):
                  A__							=   TFAutoModel.from_pretrained(UpperCAmelCase__ ,			revision='''aaaaaa''')


    def 	SCREAMING_SNAKE_CASE							(      self				:     List[Any])						->Tuple:





           '''simple docstring'''

           with self.assertRaisesRegex(
               UpperCAmelCase__ ,			'''hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin''' ,			):
                  A__							=   TFAutoModel.from_pretrained('''hf-internal-testing/config-no-model''')


    def 	SCREAMING_SNAKE_CASE							(      self				:     str)						->Optional[Any]:





           '''simple docstring'''

           with self.assertRaisesRegex(UpperCAmelCase__ ,			'''Use `from_pt=True` to load this model'''):
                  A__							=   TFAutoModel.from_pretrained('''hf-internal-testing/tiny-bert-pt-only''')





    def 	SCREAMING_SNAKE_CASE							(      self				:     List[str])						->Tuple:





           '''simple docstring'''

           A__							=   TFAutoModel.from_pretrained('''hf-internal-testing/tiny-random-bert''')
           with RequestCounter() as counter:
                  A__							=   TFAutoModel.from_pretrained('''hf-internal-testing/tiny-random-bert''')
                  self.assertEqual(counter.get_request_count ,			0)
                  self.assertEqual(counter.head_request_count ,			1)
                  self.assertEqual(counter.other_request_count ,			0)

           # With a sharded checkpoint
           A__							=   TFAutoModel.from_pretrained('''ArthurZ/tiny-random-bert-sharded''')
           with RequestCounter() as counter:
                  A__							=   TFAutoModel.from_pretrained('''ArthurZ/tiny-random-bert-sharded''')
                  self.assertEqual(counter.get_request_count ,			0)
                  self.assertEqual(counter.head_request_count ,			1)
                  self.assertEqual(counter.other_request_count ,			0)








import argparse

from tax import checkpoints

from transformers import AutoConfig, FlaxAutoModelForSeqaSeqLM







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_			,    lowercase_						)       ->	Tuple:



       """simple docstring"""

       A__							=   AutoConfig.from_pretrained(lowercase_						)
       A__							=   FlaxAutoModelForSeqaSeqLM.from_config(config=lowercase_						)
       A__							=   checkpoints.load_tax_checkpoint(lowercase_						)

       A__							=   '''wi_0''' in tax_model['''target''']['''encoder''']['''layers_0''']['''mlp''']

       if config.model_type == "t5":
              A__							=   '''SelfAttention'''
       if config.model_type == "longt5" and config.encoder_attention_type == "local":
              A__							=   '''LocalSelfAttention'''
       elif config.model_type == "longt5" and config.encoder_attention_type == "transient-global":
              A__							=   '''TransientGlobalSelfAttention'''
       else:
              raise ValueError(
                  '''Given config is expected to have `model_type=\'t5\'`, or `model_type=\'longt5` with `encoder_attention_type`'''
                  ''' attribute with a value from [\'local\', \'transient-global].'''						)

       # Encoder
       for layer_index in range(config.num_layers						):
              A__							=   f"""layers_{str(lowercase_						)}"""

              # Self-Attention
              A__							=   tax_model['''target''']['''encoder'''][layer_name]['''attention''']['''key''']['''kernel''']
              A__							=   tax_model['''target''']['''encoder'''][layer_name]['''attention''']['''out''']['''kernel''']
              A__							=   tax_model['''target''']['''encoder'''][layer_name]['''attention''']['''query''']['''kernel''']
              A__							=   tax_model['''target''']['''encoder'''][layer_name]['''attention''']['''value''']['''kernel''']

              # Global input layer norm
              if config.model_type == "longt5" and config.encoder_attention_type == "transient-global":
                     A__							=   tax_model['''target''']['''encoder'''][layer_name]['''attention''']['''T5LayerNorm_0''']['''scale''']

              # Layer Normalization
              A__							=   tax_model['''target''']['''encoder'''][layer_name]['''pre_attention_layer_norm''']['''scale''']

              if split_mlp_wi:
                     A__							=   tax_model['''target''']['''encoder'''][layer_name]['''mlp''']['''wi_0''']['''kernel''']
                     A__							=   tax_model['''target''']['''encoder'''][layer_name]['''mlp''']['''wi_1''']['''kernel''']
              else:
                     A__							=   tax_model['''target''']['''encoder'''][layer_name]['''mlp''']['''wi''']['''kernel''']

              A__							=   tax_model['''target''']['''encoder'''][layer_name]['''mlp''']['''wo''']['''kernel''']

              # Layer Normalization
              A__							=   tax_model['''target''']['''encoder'''][layer_name]['''pre_mlp_layer_norm''']['''scale''']

              # Assigning
              A__							=   flax_model.params['''encoder''']['''block'''][str(lowercase_						)]['''layer''']
              A__							=   tax_attention_key
              A__							=   tax_attention_out
              A__							=   tax_attention_query
              A__							=   tax_attention_value

              A__							=   tax_attention_layer_norm

              # Global input layer norm
              if config.model_type == "longt5" and config.encoder_attention_type == "transient-global":
                     A__							=   tax_global_layer_norm

              if split_mlp_wi:
                     A__							=   tax_mlp_wi_a
                     A__							=   tax_mlp_wi_a
              else:
                     A__							=   tax_mlp_wi

              A__							=   tax_mlp_wo
              A__							=   tax_mlp_layer_norm

              A__							=   flax_model_encoder_layer_block

       # Only for layer 0:
       A__							=   tax_model['''target''']['''encoder''']['''relpos_bias''']['''rel_embedding'''].T
       A__							=   tax_encoder_rel_embedding

       # Side/global relative position_bias + layer norm
       if config.model_type == "longt5" and config.encoder_attention_type == "transient-global":
              A__							=   tax_model['''target''']['''encoder''']['''side_relpos_bias''']['''rel_embedding'''].T
              A__							=   tax_encoder_global_rel_embedding

       # Assigning
       A__							=   tax_model['''target''']['''encoder''']['''encoder_norm''']['''scale''']
       A__							=   tax_encoder_norm

       # Decoder
       for layer_index in range(config.num_layers						):
              A__							=   f"""layers_{str(lowercase_						)}"""

              # Self-Attention
              A__							=   tax_model['''target''']['''decoder'''][layer_name]['''self_attention''']['''key''']['''kernel''']
              A__							=   tax_model['''target''']['''decoder'''][layer_name]['''self_attention''']['''out''']['''kernel''']
              A__							=   tax_model['''target''']['''decoder'''][layer_name]['''self_attention''']['''query''']['''kernel''']
              A__							=   tax_model['''target''']['''decoder'''][layer_name]['''self_attention''']['''value''']['''kernel''']

              # Layer Normalization
              A__							=   tax_model['''target''']['''decoder'''][layer_name]['''pre_self_attention_layer_norm'''][
                  '''scale'''
              ]

              # Encoder-Decoder-Attention
              A__							=   tax_model['''target''']['''decoder'''][layer_name]['''encoder_decoder_attention''']
              A__							=   tax_enc_dec_attention_module['''key''']['''kernel''']
              A__							=   tax_enc_dec_attention_module['''out''']['''kernel''']
              A__							=   tax_enc_dec_attention_module['''query''']['''kernel''']
              A__							=   tax_enc_dec_attention_module['''value''']['''kernel''']

              # Layer Normalization
              A__							=   tax_model['''target''']['''decoder'''][layer_name]['''pre_cross_attention_layer_norm''']['''scale''']

              # MLP
              if split_mlp_wi:
                     A__							=   tax_model['''target''']['''decoder'''][layer_name]['''mlp''']['''wi_0''']['''kernel''']
                     A__							=   tax_model['''target''']['''decoder'''][layer_name]['''mlp''']['''wi_1''']['''kernel''']
              else:
                     A__							=   tax_model['''target''']['''decoder'''][layer_name]['''mlp''']['''wi''']['''kernel''']

              A__							=   tax_model['''target''']['''decoder'''][layer_name]['''mlp''']['''wo''']['''kernel''']

              # Layer Normalization
              A__							=   tax_model['''target''']['''decoder'''][layer_name]['''pre_mlp_layer_norm''']['''scale''']

              # Assigning
              A__							=   flax_model.params['''decoder''']['''block'''][str(lowercase_						)]['''layer''']
              A__							=   tax_attention_key
              A__							=   tax_attention_out
              A__							=   tax_attention_query
              A__							=   tax_attention_value

              A__							=   tax_pre_attention_layer_norm

              A__							=   tax_enc_dec_attention_key
              A__							=   tax_enc_dec_attention_out
              A__							=   tax_enc_dec_attention_query
              A__							=   tax_enc_dec_attention_value

              A__							=   tax_cross_layer_norm

              if split_mlp_wi:
                     A__							=   tax_mlp_wi_a
                     A__							=   tax_mlp_wi_a
              else:
                     A__							=   tax_mlp_wi

              A__							=   tax_mlp_wo

              A__							=   txa_mlp_layer_norm

              A__							=   flax_model_decoder_layer_block

       # Decoder Normalization
       A__							=   tax_model['''target''']['''decoder''']['''decoder_norm''']['''scale''']
       A__							=   txa_decoder_norm

       # Only for layer 0:
       A__							=   tax_model['''target''']['''decoder''']['''relpos_bias''']['''rel_embedding'''].T
       A__							=   tax_decoder_rel_embedding

       # Token Embeddings
       A__							=   tax_model['''target''']['''token_embedder''']['''embedding''']
       A__							=   txa_token_embeddings

       # LM Head (only in v1.1 and LongT5 checkpoints)
       if "logits_dense" in tax_model["target"]["decoder"]:
              A__							=   tax_model['''target''']['''decoder''']['''logits_dense''']['''kernel''']

       flax_model.save_pretrained(lowercase_						)
       print('''T5X Model was sucessfully converted!'''						)


if __name__ == "__main__":
  _lowerCamelCase     :		Any								=       argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--t5x_checkpoint_path""", default=None, type=str, required=True, help="""Path the T5X checkpoint."""
  )
  parser.add_argument("""--config_name""", default=None, type=str, required=True, help="""Config name of LongT5/T5 model.""")
  parser.add_argument(
      """--flax_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output FLAX model."""
  )
  _lowerCamelCase     :		Tuple								=       parser.parse_args()
  convert_tax_checkpoint_to_flax(args.tax_checkpoint_path, args.config_name, args.flax_dump_folder_path)


