

'''simple docstring'''







from __future__ import annotations

import inspect
import unittest
from typing import List, Tuple

from transformers import RegNetConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers import TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST, TFRegNetForImageClassification, TFRegNetModel


if is_vision_available():
      from PIL import Image

      from transformers import AutoImageProcessor


class        A__   :



    def __init__(       self		,					UpperCamelCase__		,					UpperCamelCase__=3		,					UpperCamelCase__=32		,					UpperCamelCase__=3		,					UpperCamelCase__=10		,					UpperCamelCase__=[10, 20, 30, 40]		,					UpperCamelCase__=[1, 1, 2, 1]		,					UpperCamelCase__=True		,					UpperCamelCase__=True		,					UpperCamelCase__="relu"		,					UpperCamelCase__=3		,					UpperCamelCase__=None		,					)    ->   List[Any]:





           '''simple docstring'''
           A_						       =   parent
           A_						       =   batch_size
           A_						       =   image_size
           A_						       =   num_channels
           A_						       =   embeddings_size
           A_						       =   hidden_sizes
           A_						       =   depths
           A_						       =   is_training
           A_						       =   use_labels
           A_						       =   hidden_act
           A_						       =   num_labels
           A_						       =   scope
           A_						       =   len(UpperCamelCase__							)



    def        snake_case_      (       self							)    ->   Optional[int]:





           '''simple docstring'''
           A_						       =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

           A_						       =   None
           if self.use_labels:
                  A_						       =   ids_tensor([self.batch_size]		,					self.num_labels							)

           A_						       =   self.get_config()
           return config, pixel_values, labels



    def        snake_case_      (       self							)    ->   Union[str, Any]:





           '''simple docstring'''
           return RegNetConfig(
               num_channels=self.num_channels		,					embeddings_size=self.embeddings_size		,					hidden_sizes=self.hidden_sizes		,					depths=self.depths		,					hidden_act=self.hidden_act		,					num_labels=self.num_labels		,					)



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





           '''simple docstring'''
           A_						       =   TFRegNetModel(config=UpperCamelCase__							)
           A_						       =   model(UpperCamelCase__		,					training=UpperCamelCase__							)
           # expected last hidden states: B, C, H // 32, W // 32
           self.parent.assertEqual(
               result.last_hidden_state.shape		,					(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)		,					)



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)    ->   str:





           '''simple docstring'''
           A_						       =   self.num_labels
           A_						       =   TFRegNetForImageClassification(UpperCamelCase__							)
           A_						       =   model(UpperCamelCase__		,					labels=UpperCamelCase__		,					training=UpperCamelCase__							)
           self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.num_labels)							)







    def        snake_case_      (       self							)    ->   Any:





           '''simple docstring'''
           A_						       =   self.prepare_config_and_inputs()
           A_       ,      A_       ,      A_						       =   config_and_inputs
           A_						       =   {"""pixel_values""": pixel_values}
           return config, inputs_dict


@require_tf
class        A__   (   _snake_case    ,    _snake_case    ,    unittest.TestCase      ):
    lowercase						       =		(TFRegNetModel, TFRegNetForImageClassification) if is_tf_available() else ()
    lowercase						       =		(
        {"feature-extraction": TFRegNetModel, "image-classification": TFRegNetForImageClassification}
        if is_tf_available()
        else {}
    )

    lowercase						       =		False
    lowercase						       =		False
    lowercase						       =		False
    lowercase						       =		False
    lowercase						       =		False



    def        snake_case_      (       self							)    ->   str:





           '''simple docstring'''
           A_						       =   TFRegNetModelTester(self							)
           A_						       =   ConfigTester(self		,					config_class=UpperCamelCase__		,					has_text_modality=UpperCamelCase__							)



    def        snake_case_      (       self							)    ->   str:





           '''simple docstring'''
           return



    @unittest.skip(reason="""RegNet does not use inputs_embeds"""							)
    def        snake_case_      (       self							)    ->   List[str]:





           '''simple docstring'''
           pass



    @unittest.skipIf(
        not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""							)							) == 0		,					reason="""TF does not support backprop for grouped convolutions on CPU."""		,					)
    @slow
    def        snake_case_      (       self							)    ->   Dict:





           '''simple docstring'''
           super().test_keras_fit()



    @unittest.skip(reason="""RegNet does not support input and output embeddings"""							)
    def        snake_case_      (       self							)    ->   Union[str, Any]:





           '''simple docstring'''
           pass



    def        snake_case_      (       self							)    ->   Optional[int]:





           '''simple docstring'''
           A_       ,      A_						       =   self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  A_						       =   model_class(UpperCamelCase__							)
                  A_						       =   inspect.signature(model.call							)
                  # signature.parameters is an OrderedDict => so arg_names order is deterministic
                  A_						       =   [*signature.parameters.keys()]

                  A_						       =   ["""pixel_values"""]
                  self.assertListEqual(arg_names[:1]		,					UpperCamelCase__							)



    def        snake_case_      (       self							)    ->   Union[str, Any]:





           '''simple docstring'''
           A_						       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase__							)



    def        snake_case_      (       self							)    ->   int:





           '''simple docstring'''
           def check_hidden_states_output(UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							):
                  A_						       =   model_class(UpperCamelCase__							)
                  A_						       =   model(**self._prepare_for_class(UpperCamelCase__		,					UpperCamelCase__							)		,					training=UpperCamelCase__							)

                  A_						       =   outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                  A_						       =   self.model_tester.num_stages
                  self.assertEqual(len(UpperCamelCase__							)		,					expected_num_stages + 1							)

                  # RegNet's feature maps are of shape (batch_size, num_channels, height, width)
                  self.assertListEqual(
                      list(hidden_states[0].shape[-2:]							)		,					[self.model_tester.image_size // 2, self.model_tester.image_size // 2]		,					)

           A_       ,      A_						       =   self.model_tester.prepare_config_and_inputs_for_common()
           A_						       =   ["""basic""", """bottleneck"""]
           for model_class in self.all_model_classes:
                  for layer_type in layers_type:
                         A_						       =   layer_type
                         A_						       =   True
                         check_hidden_states_output(UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)

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

                         check_hidden_states_output(UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)



    def        snake_case_      (       self							)    ->   Optional[int]:





           '''simple docstring'''
           A_       ,      A_						       =   self.model_tester.prepare_config_and_inputs_for_common()

           def check_equivalence(UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__={}							):
                  A_						       =   model(UpperCamelCase__		,					return_dict=UpperCamelCase__		,					**UpperCamelCase__							)
                  A_						       =   model(UpperCamelCase__		,					return_dict=UpperCamelCase__		,					**UpperCamelCase__							).to_tuple()

                  def recursive_check(UpperCamelCase__		,					UpperCamelCase__							):
                         if isinstance(UpperCamelCase__		,					(List, Tuple)							):
                                for tuple_iterable_value, dict_iterable_value in zip(UpperCamelCase__		,					UpperCamelCase__							):
                                       recursive_check(UpperCamelCase__		,					UpperCamelCase__							)
                         elif tuple_object is None:
                                return
                         else:
                                self.assertTrue(
                                    all(tf.equal(UpperCamelCase__		,					UpperCamelCase__							)							)		,					msg=(
                                        """Tuple and dict output are not equal. Difference:"""
                                        f''' {tf.math.reduce_max(tf.abs(tuple_object - dict_object							)							)}'''
                                    )		,					)

                         recursive_check(UpperCamelCase__		,					UpperCamelCase__							)

           for model_class in self.all_model_classes:
                  A_						       =   model_class(UpperCamelCase__							)

                  A_						       =   self._prepare_for_class(UpperCamelCase__		,					UpperCamelCase__							)
                  A_						       =   self._prepare_for_class(UpperCamelCase__		,					UpperCamelCase__							)
                  check_equivalence(UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)

                  A_						       =   self._prepare_for_class(UpperCamelCase__		,					UpperCamelCase__		,					return_labels=UpperCamelCase__							)
                  A_						       =   self._prepare_for_class(UpperCamelCase__		,					UpperCamelCase__		,					return_labels=UpperCamelCase__							)
                  check_equivalence(UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__							)

                  A_						       =   self._prepare_for_class(UpperCamelCase__		,					UpperCamelCase__							)
                  A_						       =   self._prepare_for_class(UpperCamelCase__		,					UpperCamelCase__							)
                  check_equivalence(UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					{"""output_hidden_states""": True}							)

                  A_						       =   self._prepare_for_class(UpperCamelCase__		,					UpperCamelCase__		,					return_labels=UpperCamelCase__							)
                  A_						       =   self._prepare_for_class(UpperCamelCase__		,					UpperCamelCase__		,					return_labels=UpperCamelCase__							)
                  check_equivalence(UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__		,					{"""output_hidden_states""": True}							)



    def        snake_case_      (       self							)    ->   List[str]:





           '''simple docstring'''
           A_						       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCamelCase__							)







    @slow
    def        snake_case_      (       self							)    ->   Optional[int]:





           '''simple docstring'''
           for model_name in TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A_						       =   TFRegNetModel.from_pretrained(UpperCamelCase__							)
                  self.assertIsNotNone(UpperCamelCase__							)

def    UpperCAmelCase__     (   )      ->      str:
       A_						       =   Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""					)
       return image







@require_tf
@require_vision
class        A__   (   unittest.TestCase      ):



    @cached_property
    def        snake_case_      (       self							)    ->   Tuple:





           '''simple docstring'''
           return (
               AutoImageProcessor.from_pretrained(TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]							)
               if is_vision_available()
               else None
           )







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





           '''simple docstring'''
           A_						       =   TFRegNetForImageClassification.from_pretrained(TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]							)

           A_						       =   self.default_image_processor
           A_						       =   prepare_img()
           A_						       =   image_processor(images=UpperCamelCase__		,					return_tensors="""tf"""							)

           # forward pass
           A_						       =   model(**UpperCamelCase__		,					training=UpperCamelCase__							)

           # verify the logits
           A_						       =   tf.TensorShape((1, 1000)							)
           self.assertEqual(outputs.logits.shape		,					UpperCamelCase__							)

           A_						       =   tf.constant([-0.4180, -1.5051, -3.4836]							)

           tf.debugging.assert_near(outputs.logits[0, :3]		,					UpperCamelCase__		,					atol=1e-4							)



'''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"					)      ->      Dict:
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/key/kernel''']
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/out/kernel''']
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/query/kernel''']
       A_						       =   params[F'''{prefix}/layers_{i}/{layer_name}/value/kernel''']
       return k, o, q, v
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__=False					)      ->      Optional[Any]:
       if split_mlp_wi:
              A_						       =   params[F'''{prefix}/layers_{i}/mlp/wi_0/kernel''']
              A_						       =   params[F'''{prefix}/layers_{i}/mlp/wi_1/kernel''']
              A_						       =   (wi_a, wi_a)
       else:
              A_						       =   params[F'''{prefix}/layers_{i}/mlp/wi/kernel''']

       A_						       =   params[F'''{prefix}/layers_{i}/mlp/wo/kernel''']
       return wi, wo
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)      ->      Any:
       return params[F'''{prefix}/layers_{i}/{layer_name}/scale''']
def    UpperCAmelCase__     (   UpperCAmelCase__,							*, UpperCAmelCase__,							UpperCAmelCase__					)      ->      Union[str, Any]:
       A_						       =   traverse_util.flatten_dict(variables["""target"""]					)
       A_						       =   {"""/""".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
       A_						       =   """encoder/layers_0/mlp/wi_0/kernel""" in old
       print("""Split MLP:""",							UpperCAmelCase__					)

       A_						       =   collections.OrderedDict()

       # Shared embeddings.
       A_						       =   old["""token_embedder/embedding"""]

       # Encoder.
       for i in range(UpperCAmelCase__					):
              # Block i, layer 0 (Self Attention).
              A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							"""pre_attention_layer_norm"""					)
              A_       ,      A_       ,      A_       ,      A_						       =   tax_attention_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							"""attention"""					)
              A_						       =   layer_norm
              A_						       =   k.T
              A_						       =   o.T
              A_						       =   q.T
              A_						       =   v.T

              # Block i, layer 1 (MLP).
              A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							"""pre_mlp_layer_norm"""					)
              A_       ,      A_						       =   tax_mlp_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""encoder""",							UpperCAmelCase__					)
              A_						       =   layer_norm
              if split_mlp_wi:
                     A_						       =   wi[0].T
                     A_						       =   wi[1].T
              else:
                     A_						       =   wi.T
              A_						       =   wo.T

       A_						       =   old[
           """encoder/relpos_bias/rel_embedding"""
       ].T
       A_						       =   old["""encoder/encoder_norm/scale"""]

       if not is_encoder_only:
              # Decoder.
              for i in range(UpperCAmelCase__					):
                     # Block i, layer 0 (Self Attention).
                     A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""pre_self_attention_layer_norm"""					)
                     A_       ,      A_       ,      A_       ,      A_						       =   tax_attention_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""self_attention"""					)
                     A_						       =   layer_norm
                     A_						       =   k.T
                     A_						       =   o.T
                     A_						       =   q.T
                     A_						       =   v.T

                     # Block i, layer 1 (Cross Attention).
                     A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""pre_cross_attention_layer_norm"""					)
                     A_       ,      A_       ,      A_       ,      A_						       =   tax_attention_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""encoder_decoder_attention"""					)
                     A_						       =   layer_norm
                     A_						       =   k.T
                     A_						       =   o.T
                     A_						       =   q.T
                     A_						       =   v.T

                     # Block i, layer 2 (MLP).
                     A_						       =   tax_layer_norm_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							"""pre_mlp_layer_norm"""					)
                     A_       ,      A_						       =   tax_mlp_lookup(UpperCAmelCase__,							UpperCAmelCase__,							"""decoder""",							UpperCAmelCase__					)
                     A_						       =   layer_norm
                     if split_mlp_wi:
                            A_						       =   wi[0].T
                            A_						       =   wi[1].T
                     else:
                            A_						       =   wi.T
                     A_						       =   wo.T

              A_						       =   old["""decoder/decoder_norm/scale"""]
              A_						       =   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:
                     A_						       =   old["""decoder/logits_dense/kernel"""].T

       return new
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__					)      ->      List[Any]:
       A_						       =   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:
              A_						       =   state_dict["""shared.weight"""]

       if not is_encoder_only:
              if "decoder.embed_tokens.weight" not in state_dict:
                     A_						       =   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."""					)
                     A_						       =   state_dict["""shared.weight"""]

       return state_dict
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__					)      ->      Union[str, Any]:
       A_						       =   checkpoints.load_tax_checkpoint(UpperCAmelCase__					)
       A_						       =   convert_tax_to_pytorch(UpperCAmelCase__,							num_layers=config.num_layers,							is_encoder_only=UpperCAmelCase__					)
       A_						       =   make_state_dict(UpperCAmelCase__,							UpperCAmelCase__					)
       model.load_state_dict(UpperCAmelCase__,							strict=UpperCAmelCase__					)
def    UpperCAmelCase__     (   UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__,							UpperCAmelCase__ = False					)      ->      List[Any]:
       A_						       =   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:
              A_						       =   TaEncoderModel(UpperCAmelCase__					)
       else:
              A_						       =   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__":
      __lowerCamelCase									=     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
      )
      __lowerCamelCase									=     parser.parse_args()
      convert_tax_checkpoint_to_pytorch(
          args.tax_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
      )
