




import argparse
from typing import Dict

import tensorflow as tf
import torch
from tqdm import tqdm

from transformers import BigBirdPegasusConfig, BigBirdPegasusForConditionalGeneration


__a							     =					[
    # tf -> hf
    ('/', '.'),
    ('layer_', 'layers.'),
    ('kernel', 'weight'),
    ('beta', 'bias'),
    ('gamma', 'weight'),
    ('pegasus', 'model'),
]
__a							     =					[
    ('.output.dense', '.fc2'),
    ('intermediate.LayerNorm', 'final_layer_norm'),
    ('intermediate.dense', 'fc1'),
]

__a							     =					(
    INIT_COMMON
    + [
        ('attention.self.LayerNorm', 'self_attn_layer_norm'),
        ('attention.output.dense', 'self_attn.out_proj'),
        ('attention.self', 'self_attn'),
        ('attention.encdec.LayerNorm', 'encoder_attn_layer_norm'),
        ('attention.encdec_output.dense', 'encoder_attn.out_proj'),
        ('attention.encdec', 'encoder_attn'),
        ('key', 'k_proj'),
        ('value', 'v_proj'),
        ('query', 'q_proj'),
        ('decoder.LayerNorm', 'decoder.layernorm_embedding'),
    ]
    + END_COMMON
)

__a							     =					(
    INIT_COMMON
    + [
        ('embeddings.word_embeddings', 'shared.weight'),
        ('embeddings.position_embeddings', 'embed_positions.weight'),
        ('attention.self.LayerNorm', 'self_attn_layer_norm'),
        ('attention.output.dense', 'self_attn.output'),
        ('attention.self', 'self_attn.self'),
        ('encoder.LayerNorm', 'encoder.layernorm_embedding'),
    ]
    + END_COMMON
)

__a							     =					[
    'encdec/key/bias',
    'encdec/query/bias',
    'encdec/value/bias',
    'self/key/bias',
    'self/query/bias',
    'self/value/bias',
    'encdec_output/dense/bias',
    'attention/output/dense/bias',
]







def lowerCamelCase__  (     _lowercase					,		_lowercase ):

     '''simple docstring'''


     for tf_name, hf_name in patterns:
          UpperCAmelCase_     :				int  =     k.replace(_lowercase					,		_lowercase )
     return k







def lowerCamelCase__  (     _lowercase					,		_lowercase ):

     '''simple docstring'''


     UpperCAmelCase_     :				Union[str, Any]  =     BigBirdPegasusConfig(**_lowercase )
     UpperCAmelCase_     :				Union[str, Any]  =     BigBirdPegasusForConditionalGeneration(_lowercase )
     UpperCAmelCase_     :				Union[str, Any]  =     torch_model.state_dict()
     UpperCAmelCase_     :				Dict  =     {}

     # separating decoder weights
     UpperCAmelCase_     :				Any  =     {k: tf_weights[k] for k in tf_weights if k.startswith('''pegasus/decoder''' )}
     UpperCAmelCase_     :				Optional[Any]  =     {k: tf_weights[k] for k in tf_weights if not k.startswith('''pegasus/decoder''' )}

     for k, v in tqdm(decoder_weights.items()					,		'''tf -> hf conversion''' ):
          UpperCAmelCase_     :				Optional[Any]  =     [k.endswith(_lowercase ) for ending in KEYS_TO_IGNORE]
          if any(_lowercase ):
               continue
          UpperCAmelCase_     :				Optional[int]  =     DECODER_PATTERNS
          UpperCAmelCase_     :				Union[str, Any]  =     rename_state_dict_key(_lowercase					,		_lowercase )
          if new_k not in state_dict:
               raise ValueError(f'''could not find new key {new_k} in state dict. (converted from {k})''' )
          if any(True if i in k else False for i in ['''dense''', '''query''', '''key''', '''value'''] ):
               UpperCAmelCase_     :				int  =     v.T
          UpperCAmelCase_     :				List[Any]  =     torch.from_numpy(_lowercase )
          assert v.shape == state_dict[new_k].shape, f'''{new_k}, {k}, {v.shape}, {state_dict[new_k].shape}'''

     for k, v in tqdm(remaining_weights.items()					,		'''tf -> hf conversion''' ):
          UpperCAmelCase_     :				Optional[int]  =     [k.endswith(_lowercase ) for ending in KEYS_TO_IGNORE]
          if any(_lowercase ):
               continue
          UpperCAmelCase_     :				str  =     REMAINING_PATTERNS
          UpperCAmelCase_     :				Union[str, Any]  =     rename_state_dict_key(_lowercase					,		_lowercase )
          if new_k not in state_dict and k != "pegasus/embeddings/position_embeddings":
               raise ValueError(f'''could not find new key {new_k} in state dict. (converted from {k})''' )
          if any(True if i in k else False for i in ['''dense''', '''query''', '''key''', '''value'''] ):
               UpperCAmelCase_     :				int  =     v.T
          UpperCAmelCase_     :				Tuple  =     torch.from_numpy(_lowercase )
          if k != "pegasus/embeddings/position_embeddings":
               assert v.shape == state_dict[new_k].shape, f'''{new_k}, {k}, {v.shape}, {state_dict[new_k].shape}'''

     UpperCAmelCase_     :				Tuple  =     mapping['''model.embed_positions.weight''']
     UpperCAmelCase_     :				str  =     mapping.pop('''model.embed_positions.weight''' )
     UpperCAmelCase_,       UpperCAmelCase_     :				List[Any]  =     torch_model.load_state_dict(_lowercase					,		strict=_lowercase )
     UpperCAmelCase_     :				str  =     [
         k
         for k in missing
         if k
         not in [
             '''final_logits_bias''',
             '''model.encoder.embed_tokens.weight''',
             '''model.decoder.embed_tokens.weight''',
             '''lm_head.weight''',
         ]
     ]
     assert unexpected_missing == [], f'''no matches found for the following torch keys {unexpected_missing}'''
     assert extra == [], f'''no matches found for the following tf keys {extra}'''
     return torch_model







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


     UpperCAmelCase_     :				List[str]  =     tf.train.list_variables(_lowercase )
     UpperCAmelCase_     :				List[Any]  =     {}
     UpperCAmelCase_     :				List[str]  =     ['''global_step''']
     for name, shape in tqdm(_lowercase					,		desc='''converting tf checkpoint to dict''' ):
          UpperCAmelCase_     :				Any  =     any(pat in name for pat in ignore_name )
          if skip_key:
               continue
          UpperCAmelCase_     :				Tuple  =     tf.train.load_variable(_lowercase					,		_lowercase )
          UpperCAmelCase_     :				List[Any]  =     array
     return tf_weights







def lowerCamelCase__  (     _lowercase					,		_lowercase					,		_lowercase ):

     '''simple docstring'''


     UpperCAmelCase_     :				Dict  =     get_tf_weights_as_numpy(_lowercase )
     UpperCAmelCase_     :				Any  =     convert_bigbird_pegasus(_lowercase					,		_lowercase )
     torch_model.save_pretrained(_lowercase )


if __name__ == "__main__":
     __a							     =					argparse.ArgumentParser()
     parser.add_argument('--tf_ckpt_path', type=str, help='passed to tf.train.list_variables')
     parser.add_argument('--save_dir', default=None, type=str, help='Path to the output PyTorch model.')
     __a							     =					parser.parse_args()
     __a							     =					{}
     convert_bigbird_pegasus_ckpt_to_pytorch(args.tf_ckpt_path, args.save_dir, config_update=config_update)






import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
    import jax
    import jax.numpy as jnp

    from transformers.modeling_flax_pytorch_utils import (
        convert_pytorch_state_dict_to_flax,
        load_flax_weights_in_pytorch_model,
    )
    from transformers.models.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
    import torch

class   lowerCAmelCase				:
 def __init__( self    : List[str]     ,	UpperCAmelCase    : List[Any]     ,	UpperCAmelCase    : str=14     ,	UpperCAmelCase    : Any=7     ,	UpperCAmelCase    : Dict=True     ,	UpperCAmelCase    : List[str]=True     ,	UpperCAmelCase    : Any=False     ,	UpperCAmelCase    : Tuple=True     ,	UpperCAmelCase    : int=99     ,	UpperCAmelCase    : str=32     ,	UpperCAmelCase    : int=4     ,	UpperCAmelCase    : Optional[Any]=4     ,	UpperCAmelCase    : Any=4     ,	UpperCAmelCase    : Optional[int]=37     ,	UpperCAmelCase    : Union[str, Any]="gelu"     ,	UpperCAmelCase    : List[Any]=0.1     ,	UpperCAmelCase    : Dict=0.1     ,	UpperCAmelCase    : Optional[int]=512     ,	UpperCAmelCase    : List[Any]=0.0_2     ,	)     ->      List[Any]:
     lowerCamelCase__       : Tuple   					=						parent
     lowerCamelCase__       : Tuple   					=						batch_size
     lowerCamelCase__       : Optional[Any]   					=						seq_length
     lowerCamelCase__       : str   					=						is_training
     lowerCamelCase__       : Any   					=						use_input_mask
     lowerCamelCase__       : Optional[Any]   					=						use_token_type_ids
     lowerCamelCase__       : Optional[Any]   					=						use_labels
     lowerCamelCase__       : Optional[Any]   					=						vocab_size
     lowerCamelCase__       : Optional[int]   					=						hidden_size
     lowerCamelCase__       : Optional[Any]   					=						rotary_dim
     lowerCamelCase__       : Optional[Any]   					=						num_hidden_layers
     lowerCamelCase__       : Optional[Any]   					=						num_attention_heads
     lowerCamelCase__       : Tuple   					=						intermediate_size
     lowerCamelCase__       : Union[str, Any]   					=						hidden_act
     lowerCamelCase__       : List[Any]   					=						hidden_dropout_prob
     lowerCamelCase__       : int   					=						attention_probs_dropout_prob
     lowerCamelCase__       : Dict   					=						max_position_embeddings
     lowerCamelCase__       : Tuple   					=						initializer_range
     lowerCamelCase__       : Optional[int]   					=						None
     lowerCamelCase__       : int   					=						vocab_size - 1
     lowerCamelCase__       : str   					=						vocab_size - 1
     lowerCamelCase__       : str   					=						vocab_size - 1
 def 	A_      ( self    : str      )     ->      int:
     lowerCamelCase__       : Optional[int]   					=						ids_tensor([self.batch_size, self.seq_length]     ,	self.vocab_size      )

     lowerCamelCase__       : Optional[Any]   					=						None
     if self.use_input_mask:
         lowerCamelCase__       : List[str]   					=						random_attention_mask([self.batch_size, self.seq_length]      )

     lowerCamelCase__       : Tuple   					=						GPTJConfig(
         vocab_size=self.vocab_size     ,	n_embd=self.hidden_size     ,	n_layer=self.num_hidden_layers     ,	n_head=self.num_attention_heads     ,	n_positions=self.max_position_embeddings     ,	use_cache=UpperCAmelCase     ,	bos_token_id=self.bos_token_id     ,	eos_token_id=self.eos_token_id     ,	pad_token_id=self.pad_token_id     ,	rotary_dim=self.rotary_dim     ,	)

     return (config, input_ids, input_mask)
 def 	A_      ( self    : Optional[Any]      )     ->      Union[str, Any]:
     lowerCamelCase__       : Tuple   					=						self.prepare_config_and_inputs()
     lowerCamelCase__					,				lowerCamelCase__					,				lowerCamelCase__       : List[Any]   					=						config_and_inputs
     lowerCamelCase__       : List[Any]   					=						{'input_ids': input_ids, 'attention_mask': attention_mask}
     return config, inputs_dict
 def 	A_      ( self    : int     ,	UpperCAmelCase    : Optional[Any]     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : List[str]     ,	UpperCAmelCase    : Optional[Any]      )     ->      Tuple:
     lowerCamelCase__       : Tuple   					=						20
     lowerCamelCase__       : Dict   					=						model_class_name(UpperCAmelCase      )

     lowerCamelCase__       : Dict   					=						model.init_cache(input_ids.shape[0]     ,	UpperCAmelCase      )
     lowerCamelCase__       : int   					=						jnp.ones((input_ids.shape[0], max_decoder_length)     ,	dtype='i4'      )

     lowerCamelCase__       : List[str]   					=						jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1      )[None, :]     ,	(input_ids.shape[0], input_ids.shape[-1] - 1)      )
     lowerCamelCase__       : Optional[int]   					=						model(
         input_ids[:, :-1]     ,	attention_mask=UpperCAmelCase     ,	past_key_values=UpperCAmelCase     ,	position_ids=UpperCAmelCase     ,	)

     lowerCamelCase__       : List[str]   					=						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]     ,	dtype='i4'      )
     lowerCamelCase__       : List[str]   					=						model(
         input_ids[:, -1:]     ,	attention_mask=UpperCAmelCase     ,	past_key_values=outputs_cache.past_key_values     ,	position_ids=UpperCAmelCase     ,	)

     lowerCamelCase__       : List[str]   					=						model(UpperCAmelCase      )

     lowerCamelCase__       : Optional[Any]   					=						np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
     self.parent.assertTrue(diff < 1e-3     ,	msg=F"""Max diff is {diff}"""      )







 def 	A_      ( self    : Union[str, Any]     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : Any     ,	UpperCAmelCase    : List[Any]     ,	UpperCAmelCase    : List[str]      )     ->      Optional[Any]:
     lowerCamelCase__       : Any   					=						20
     lowerCamelCase__       : Any   					=						model_class_name(UpperCAmelCase      )

     lowerCamelCase__       : List[str]   					=						jnp.concatenate(
         [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])      )]     ,	axis=-1     ,	)

     lowerCamelCase__       : str   					=						model.init_cache(input_ids.shape[0]     ,	UpperCAmelCase      )
     lowerCamelCase__       : Dict   					=						jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1      )[None, :]     ,	(input_ids.shape[0], input_ids.shape[-1] - 1)      )

     lowerCamelCase__       : List[Any]   					=						model(
         input_ids[:, :-1]     ,	attention_mask=UpperCAmelCase     ,	past_key_values=UpperCAmelCase     ,	position_ids=UpperCAmelCase     ,	)
     lowerCamelCase__       : Optional[Any]   					=						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]     ,	dtype='i4'      )
     lowerCamelCase__       : List[Any]   					=						model(
         input_ids[:, -1:]     ,	past_key_values=outputs_cache.past_key_values     ,	attention_mask=UpperCAmelCase     ,	position_ids=UpperCAmelCase     ,	)

     lowerCamelCase__       : List[Any]   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase      )

     lowerCamelCase__       : Dict   					=						np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
     self.parent.assertTrue(diff < 1e-3     ,	msg=F"""Max diff is {diff}"""      )


@require_flax
class   lowerCAmelCase				(   __UpperCamelCase,    __UpperCamelCase,    unittest.TestCase     ):
 UpperCAmelCase__						=						(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
 UpperCAmelCase__						=						(FlaxGPTJForCausalLM,) if is_flax_available() else ()
 def 	A_      ( self    : Optional[Any]      )     ->      Tuple:
     lowerCamelCase__       : Optional[Any]   					=						FlaxGPTJModelTester(self      )
 def 	A_      ( self    : Tuple      )     ->      Union[str, Any]:
     for model_class_name in self.all_model_classes:
         lowerCamelCase__					,				lowerCamelCase__					,				lowerCamelCase__       : Optional[int]   					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.check_use_cache_forward(UpperCAmelCase     ,	UpperCAmelCase     ,	UpperCAmelCase     ,	UpperCAmelCase      )
 def 	A_      ( self    : Union[str, Any]      )     ->      Tuple:
     for model_class_name in self.all_model_classes:
         lowerCamelCase__					,				lowerCamelCase__					,				lowerCamelCase__       : str   					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.check_use_cache_forward_with_attn_mask(
             UpperCAmelCase     ,	UpperCAmelCase     ,	UpperCAmelCase     ,	UpperCAmelCase      )
 @tooslow
 def 	A_      ( self    : Dict      )     ->      Tuple:
     lowerCamelCase__       : str   					=						GPTaTokenizer.from_pretrained('gpt2'     ,	pad_token='<|endoftext|>'     ,	padding_side='left'      )
     lowerCamelCase__       : Dict   					=						tokenizer(['Hello this is a long string', 'Hey']     ,	return_tensors='np'     ,	padding=UpperCAmelCase     ,	truncation=UpperCAmelCase      )

     lowerCamelCase__       : str   					=						FlaxGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'      )
     lowerCamelCase__       : Optional[int]   					=						False
     lowerCamelCase__       : Optional[int]   					=						model.config.eos_token_id

     lowerCamelCase__       : Union[str, Any]   					=						jax.jit(model.generate      )

     lowerCamelCase__       : Optional[Any]   					=						jit_generate(
         inputs['input_ids']     ,	attention_mask=inputs['attention_mask']     ,	pad_token_id=tokenizer.pad_token_id      ).sequences

     lowerCamelCase__       : Optional[int]   					=						tokenizer.batch_decode(UpperCAmelCase     ,	skip_special_tokens=UpperCAmelCase      )

     lowerCamelCase__       : int   					=						[
         'Hello this is a long string of text.\n\nI\'m trying to get the text of the',
         'Hey, I\'m a little late to the party. I\'m going to',
     ]

     self.assertListEqual(UpperCAmelCase     ,	UpperCAmelCase      )
 @is_pt_flax_cross_test
 def 	A_      ( self    : List[str]      )     ->      Union[str, Any]:
     lowerCamelCase__					,				lowerCamelCase__       : Optional[int]   					=						self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         with self.subTest(model_class.__name__      ):
             # prepare inputs
             lowerCamelCase__       : Dict   					=						self._prepare_for_class(UpperCAmelCase     ,	UpperCAmelCase      )
             lowerCamelCase__       : Optional[int]   					=						{k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

             # load corresponding PyTorch class
             lowerCamelCase__       : Any   					=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
             lowerCamelCase__       : Union[str, Any]   					=						getattr(UpperCAmelCase     ,	UpperCAmelCase      )

             lowerCamelCase__					,				lowerCamelCase__       : Tuple   					=						pt_inputs['input_ids'].shape
             lowerCamelCase__       : Tuple   					=						np.random.randint(0     ,	seq_length - 1     ,	size=(batch_size,)      )
             for batch_idx, start_index in enumerate(UpperCAmelCase      ):
                 lowerCamelCase__       : List[str]   					=						0
                 lowerCamelCase__       : Optional[int]   					=						1
                 lowerCamelCase__       : Any   					=						0
                 lowerCamelCase__       : Any   					=						1
             lowerCamelCase__       : Dict   					=						pt_model_class(UpperCAmelCase      ).eval()
             lowerCamelCase__       : str   					=						model_class(UpperCAmelCase     ,	dtype=jnp.floataa      )

             lowerCamelCase__       : str   					=						convert_pytorch_state_dict_to_flax(pt_model.state_dict()     ,	UpperCAmelCase      )
             lowerCamelCase__       : Tuple   					=						fx_state

             with torch.no_grad():
                 lowerCamelCase__       : List[str]   					=						pt_model(**UpperCAmelCase      ).to_tuple()

             lowerCamelCase__       : Dict   					=						fx_model(**UpperCAmelCase      ).to_tuple()
             self.assertEqual(len(UpperCAmelCase      )     ,	len(UpperCAmelCase      )     ,	'Output lengths differ between Flax and PyTorch'      )
             for fx_output, pt_output in zip(UpperCAmelCase     ,	UpperCAmelCase      ):
                 self.assert_almost_equals(fx_output[:, -1]     ,	pt_output[:, -1].numpy()     ,	4e-2      )

             with tempfile.TemporaryDirectory() as tmpdirname:
                 pt_model.save_pretrained(UpperCAmelCase      )
                 lowerCamelCase__       : Tuple   					=						model_class.from_pretrained(UpperCAmelCase     ,	from_pt=UpperCAmelCase      )

             lowerCamelCase__       : int   					=						fx_model_loaded(**UpperCAmelCase      ).to_tuple()
             self.assertEqual(
                 len(UpperCAmelCase      )     ,	len(UpperCAmelCase      )     ,	'Output lengths differ between Flax and PyTorch'      )
             for fx_output_loaded, pt_output in zip(UpperCAmelCase     ,	UpperCAmelCase      ):
                 self.assert_almost_equals(fx_output_loaded[:, -1]     ,	pt_output[:, -1].numpy()     ,	4e-2      )
 @is_pt_flax_cross_test
 def 	A_      ( self    : Union[str, Any]      )     ->      int:
     lowerCamelCase__					,				lowerCamelCase__       : Dict   					=						self.model_tester.prepare_config_and_inputs_for_common()
     for model_class in self.all_model_classes:
         with self.subTest(model_class.__name__      ):
             # prepare inputs
             lowerCamelCase__       : Tuple   					=						self._prepare_for_class(UpperCAmelCase     ,	UpperCAmelCase      )
             lowerCamelCase__       : Tuple   					=						{k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

             # load corresponding PyTorch class
             lowerCamelCase__       : Union[str, Any]   					=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
             lowerCamelCase__       : int   					=						getattr(UpperCAmelCase     ,	UpperCAmelCase      )

             lowerCamelCase__       : List[Any]   					=						pt_model_class(UpperCAmelCase      ).eval()
             lowerCamelCase__       : str   					=						model_class(UpperCAmelCase     ,	dtype=jnp.floataa      )

             lowerCamelCase__       : List[str]   					=						load_flax_weights_in_pytorch_model(UpperCAmelCase     ,	fx_model.params      )
             lowerCamelCase__					,				lowerCamelCase__       : Union[str, Any]   					=						pt_inputs['input_ids'].shape
             lowerCamelCase__       : Dict   					=						np.random.randint(0     ,	seq_length - 1     ,	size=(batch_size,)      )
             for batch_idx, start_index in enumerate(UpperCAmelCase      ):
                 lowerCamelCase__       : List[Any]   					=						0
                 lowerCamelCase__       : Tuple   					=						1
                 lowerCamelCase__       : Optional[Any]   					=						0
                 lowerCamelCase__       : List[str]   					=						1

             # make sure weights are tied in PyTorch
             pt_model.tie_weights()

             with torch.no_grad():
                 lowerCamelCase__       : Any   					=						pt_model(**UpperCAmelCase      ).to_tuple()

             lowerCamelCase__       : List[Any]   					=						fx_model(**UpperCAmelCase      ).to_tuple()
             self.assertEqual(len(UpperCAmelCase      )     ,	len(UpperCAmelCase      )     ,	'Output lengths differ between Flax and PyTorch'      )
             for fx_output, pt_output in zip(UpperCAmelCase     ,	UpperCAmelCase      ):
                 self.assert_almost_equals(fx_output[:, -1]     ,	pt_output[:, -1].numpy()     ,	4e-2      )

             with tempfile.TemporaryDirectory() as tmpdirname:
                 fx_model.save_pretrained(UpperCAmelCase      )
                 lowerCamelCase__       : Tuple   					=						pt_model_class.from_pretrained(UpperCAmelCase     ,	from_flax=UpperCAmelCase      )

             with torch.no_grad():
                 lowerCamelCase__       : Optional[int]   					=						pt_model_loaded(**UpperCAmelCase      ).to_tuple()

             self.assertEqual(
                 len(UpperCAmelCase      )     ,	len(UpperCAmelCase      )     ,	'Output lengths differ between Flax and PyTorch'      )
             for fx_output, pt_output in zip(UpperCAmelCase     ,	UpperCAmelCase      ):
                 self.assert_almost_equals(fx_output[:, -1]     ,	pt_output[:, -1].numpy()     ,	4e-2      )







 @tooslow
 def 	A_      ( self    : List[str]      )     ->      List[Any]:
     for model_class_name in self.all_model_classes:
         lowerCamelCase__       : Union[str, Any]   					=						model_class_name.from_pretrained('EleutherAI/gpt-j-6B'      )
         lowerCamelCase__       : Any   					=						model(np.ones((1, 1)      )      )
         self.assertIsNotNone(UpperCAmelCase      )





