


'''simple docstring'''





from __future__ import annotations

import unittest
import warnings

from transformers import AutoTokenizer, MarianConfig, MarianTokenizer, TranslationPipeline, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
from transformers.utils import cached_property

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


if is_tf_available():
  import tensorflow as tf

  from transformers import TFAutoModelForSeqaSeqLM, TFMarianModel, TFMarianMTModel







@require_tf
class    lowerCamelCase__     :
     SCREAMING_SNAKE_CASE__       =    MarianConfig
     SCREAMING_SNAKE_CASE__       =    {}
     SCREAMING_SNAKE_CASE__       =    '''gelu'''





     def __init__(self	,  UpperCAmelCase	,  UpperCAmelCase=1_3	,  UpperCAmelCase=7	,  UpperCAmelCase=True	,  UpperCAmelCase=False	,  UpperCAmelCase=9_9	,  UpperCAmelCase=3_2	,  UpperCAmelCase=2	,  UpperCAmelCase=4	,  UpperCAmelCase=3_7	,  UpperCAmelCase=0.1	,  UpperCAmelCase=0.1	,  UpperCAmelCase=2_0	,  UpperCAmelCase=2	,  UpperCAmelCase=1	,  UpperCAmelCase=0	,  )  ->     Union[str, Any]:
      _lowercase							   =parent
      _lowercase							   =batch_size
      _lowercase							   =seq_length
      _lowercase							   =is_training
      _lowercase							   =use_labels
      _lowercase							   =vocab_size
      _lowercase							   =hidden_size
      _lowercase							   =num_hidden_layers
      _lowercase							   =num_attention_heads
      _lowercase							   =intermediate_size

      _lowercase							   =hidden_dropout_prob
      _lowercase							   =attention_probs_dropout_prob
      _lowercase							   =max_position_embeddings
      _lowercase							   =eos_token_id
      _lowercase							   =pad_token_id
      _lowercase							   =bos_token_id





     def   __A      (self )  ->     Dict:
      _lowercase							   =ids_tensor([self.batch_size, self.seq_length - 1]	,  self.vocab_size )
      _lowercase							   =tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size )	,  1 )
      _lowercase							   =tf.concat([input_ids, eos_tensor]	,  axis=1 )

      _lowercase							   =ids_tensor([self.batch_size, self.seq_length]	,  self.vocab_size )

      _lowercase							   =self.config_cls(
          vocab_size=self.vocab_size	,  d_model=self.hidden_size	,  encoder_layers=self.num_hidden_layers	,  decoder_layers=self.num_hidden_layers	,  encoder_attention_heads=self.num_attention_heads	,  decoder_attention_heads=self.num_attention_heads	,  encoder_ffn_dim=self.intermediate_size	,  decoder_ffn_dim=self.intermediate_size	,  dropout=self.hidden_dropout_prob	,  attention_dropout=self.attention_probs_dropout_prob	,  max_position_embeddings=self.max_position_embeddings	,  eos_token_ids=[2]	,  bos_token_id=self.bos_token_id	,  pad_token_id=self.pad_token_id	,  decoder_start_token_id=self.pad_token_id	,  **self.config_updates	,  )
      _lowercase							   =prepare_marian_inputs_dict(UpperCamelCase__	,  UpperCamelCase__	,  UpperCamelCase__ )
      return config, inputs_dict





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase )  ->     Optional[int]:
      _lowercase							   =TFMarianModel(config=UpperCamelCase__ ).get_decoder()
      _lowercase							   =inputs_dict['''input_ids''']

      _lowercase							   =input_ids[:1, :]
      _lowercase							   =inputs_dict['''attention_mask'''][:1, :]
      _lowercase							   =inputs_dict['''head_mask''']
      _lowercase							   =1

      # first forward pass
      _lowercase							   =model(UpperCamelCase__	,  attention_mask=UpperCamelCase__	,  head_mask=UpperCamelCase__	,  use_cache=UpperCamelCase__ )

      _lowercase							   =outputs.to_tuple()

      # create hypothetical next token and extent to next_input_ids
      _lowercase							   =ids_tensor((self.batch_size, 3)	,  config.vocab_size )
      _lowercase							   =tf.cast(ids_tensor((self.batch_size, 3)	,  2 )	,  tf.inta )

      # append to next input_ids and
      _lowercase							   =tf.concat([input_ids, next_tokens]	,  axis=-1 )
      _lowercase							   =tf.concat([attention_mask, next_attn_mask]	,  axis=-1 )

      _lowercase							   =model(UpperCamelCase__	,  attention_mask=UpperCamelCase__ )[0]
      _lowercase							   =model(UpperCamelCase__	,  attention_mask=UpperCamelCase__	,  past_key_values=UpperCamelCase__ )[0]

      self.parent.assertEqual(next_tokens.shape[1]	,  output_from_past.shape[1] )

      # select random slice
      _lowercase							   =int(ids_tensor((1,)	,  output_from_past.shape[-1] ) )
      _lowercase							   =output_from_no_past[:, -3:, random_slice_idx]
      _lowercase							   =output_from_past[:, :, random_slice_idx]

      # test that outputs are equal for slice
      tf.debugging.assert_near(UpperCamelCase__	,  UpperCamelCase__	,  rtol=1e-3 )



def        UpperCAmelCase_      (							__snake_case  ,		__snake_case  ,		__snake_case  ,		__snake_case=None  ,		__snake_case=None  ,		__snake_case=None  ,		__snake_case=None  ,		__snake_case=None  ,		)  ->    Optional[Any]:







 """simple docstring"""


 if attention_mask is None:
  _lowercase							   =tf.cast(tf.math.not_equal(SCREAMING_SNAKE_CASE__  ,		config.pad_token_id		)  ,		tf.inta		)
 if decoder_attention_mask is None:
  _lowercase							   =tf.concat(
      [
          tf.ones(decoder_input_ids[:, :1].shape  ,		dtype=tf.inta		),
          tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]  ,		config.pad_token_id		)  ,		tf.inta		),
      ]  ,		axis=-1  ,		)
 if head_mask is None:
  _lowercase							   =tf.ones((config.encoder_layers, config.encoder_attention_heads)		)
 if decoder_head_mask is None:
  _lowercase							   =tf.ones((config.decoder_layers, config.decoder_attention_heads)		)
 if cross_attn_head_mask is None:
  _lowercase							   =tf.ones((config.decoder_layers, config.decoder_attention_heads)		)
 return {
     "input_ids": input_ids,
     "decoder_input_ids": decoder_input_ids,
     "attention_mask": attention_mask,
     "decoder_attention_mask": decoder_attention_mask,
     "head_mask": head_mask,
     "decoder_head_mask": decoder_head_mask,
     "cross_attn_head_mask": cross_attn_head_mask,
 }







@require_tf
class    lowerCamelCase__     ( __SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							unittest.TestCase):
     SCREAMING_SNAKE_CASE__       =    (TFMarianMTModel, TFMarianModel) if is_tf_available() else ()
     SCREAMING_SNAKE_CASE__       =    (TFMarianMTModel,) if is_tf_available() else ()
     SCREAMING_SNAKE_CASE__       =    (
         {
             '''conversational''': TFMarianMTModel,
             '''feature-extraction''': TFMarianModel,
             '''summarization''': TFMarianMTModel,
             '''text2text-generation''': TFMarianMTModel,
             '''translation''': TFMarianMTModel,
         }
         if is_tf_available()
         else {}
     )
     SCREAMING_SNAKE_CASE__       =    True
     SCREAMING_SNAKE_CASE__       =    False
     SCREAMING_SNAKE_CASE__       =    False





     def   __A      (self )  ->     List[Any]:
      _lowercase							   =TFMarianModelTester(self )
      _lowercase							   =ConfigTester(self	,  config_class=UpperCamelCase__ )





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





     def   __A      (self )  ->     Optional[int]:
      _lowercase							   =self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.check_decoder_model_past_large_inputs(*UpperCamelCase__ )







@require_tf
class    lowerCamelCase__     ( unittest.TestCase):
     SCREAMING_SNAKE_CASE__       =    1000  # show more chars for failing integration tests





     @classmethod
     def   __A      (cls )  ->     None:
      _lowercase							   =f"Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}"
      return cls





     @cached_property
     def   __A      (self )  ->     MarianTokenizer:
      return AutoTokenizer.from_pretrained(self.model_name )





     @property
     def   __A      (self )  ->     int:
      return self.tokenizer.eos_token_id





     @cached_property
     def   __A      (self )  ->     Tuple:
      warnings.simplefilter('''error''' )
      _lowercase							   =TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name )
      assert isinstance(UpperCamelCase__	,  UpperCamelCase__ )
      _lowercase							   =model.config
      self.assertListEqual(c.bad_words_ids	,  [[c.pad_token_id]] )
      self.assertEqual(c.max_length	,  5_1_2 )
      self.assertEqual(c.decoder_start_token_id	,  c.pad_token_id )
      return model





     def   __A      (self	,  **UpperCAmelCase )  ->     int:
      _lowercase							   =self.translate_src_text(**UpperCamelCase__ )
      self.assertListEqual(self.expected_text	,  UpperCamelCase__ )





     def   __A      (self	,  **UpperCAmelCase )  ->     Optional[int]:
      _lowercase							   =self.tokenizer(self.src_text	,  **UpperCamelCase__	,  padding=UpperCamelCase__	,  return_tensors='''tf''' )
      _lowercase							   =self.model.generate(
          model_inputs.input_ids	,  attention_mask=model_inputs.attention_mask	,  num_beams=2	,  max_length=1_2_8 )
      _lowercase							   =self.tokenizer.batch_decode(generated_ids.numpy()	,  skip_special_tokens=UpperCamelCase__ )
      return generated_words







@require_sentencepiece
@require_tokenizers
@require_tf
class    lowerCamelCase__     ( __SCREAMING_SNAKE_CASE):
     SCREAMING_SNAKE_CASE__       =    '''mt'''
     SCREAMING_SNAKE_CASE__       =    '''en'''
     SCREAMING_SNAKE_CASE__       =    ['''Billi messu b\'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem.''']
     SCREAMING_SNAKE_CASE__       =    ['''Touching gently, Jesus healed a man who was affected by the sad disease of leprosy.''']





     @unittest.skip('''Skipping until #12647 is resolved.''' )
     @slow
     def   __A      (self )  ->     List[str]:
      self._assert_generated_batch_equal_expected()







@require_sentencepiece
@require_tokenizers
@require_tf
class    lowerCamelCase__     ( __SCREAMING_SNAKE_CASE):
     SCREAMING_SNAKE_CASE__       =    '''en'''
     SCREAMING_SNAKE_CASE__       =    '''zh'''
     SCREAMING_SNAKE_CASE__       =    ['''My name is Wolfgang and I live in Berlin''']
     SCREAMING_SNAKE_CASE__       =    ['''我叫沃尔夫冈 我住在柏林''']





     @unittest.skip('''Skipping until #12647 is resolved.''' )
     @slow
     def   __A      (self )  ->     List[Any]:
      self._assert_generated_batch_equal_expected()







@require_sentencepiece
@require_tokenizers
@require_tf
class    lowerCamelCase__     ( __SCREAMING_SNAKE_CASE):
     SCREAMING_SNAKE_CASE__       =    '''en'''
     SCREAMING_SNAKE_CASE__       =    '''ROMANCE'''
     SCREAMING_SNAKE_CASE__       =    [
         '''>>fr<< Don\'t spend so much time watching TV.''',
         '''>>pt<< Your message has been sent.''',
         '''>>es<< He\'s two years older than me.''',
     ]
     SCREAMING_SNAKE_CASE__       =    [
         '''Ne passez pas autant de temps à regarder la télé.''',
         '''A sua mensagem foi enviada.''',
         '''Es dos años más viejo que yo.''',
     ]





     @unittest.skip('''Skipping until #12647 is resolved.''' )
     @slow
     def   __A      (self )  ->     int:
      self._assert_generated_batch_equal_expected()





     @unittest.skip('''Skipping until #12647 is resolved.''' )
     @slow
     def   __A      (self )  ->     Any:
      _lowercase							   =TranslationPipeline(self.model	,  self.tokenizer	,  framework='''tf''' )
      _lowercase							   =pipeline(self.src_text )
      self.assertEqual(self.expected_text	,  [x['''translation_text'''] for x in output] )





import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, 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.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
  import torch







class    lowerCamelCase__     :





     def __init__(self	,  UpperCAmelCase	,  UpperCAmelCase=1_4	,  UpperCAmelCase=7	,  UpperCAmelCase=True	,  UpperCAmelCase=True	,  UpperCAmelCase=False	,  UpperCAmelCase=True	,  UpperCAmelCase=9_9	,  UpperCAmelCase=3_2	,  UpperCAmelCase=5	,  UpperCAmelCase=4	,  UpperCAmelCase=3_7	,  UpperCAmelCase="gelu"	,  UpperCAmelCase=0.1	,  UpperCAmelCase=0.1	,  UpperCAmelCase=5_1_2	,  UpperCAmelCase=0.02	,  )  ->     Any:
      _lowercase							   =parent
      _lowercase							   =batch_size
      _lowercase							   =seq_length
      _lowercase							   =is_training
      _lowercase							   =use_input_mask
      _lowercase							   =use_token_type_ids
      _lowercase							   =use_labels
      _lowercase							   =vocab_size
      _lowercase							   =hidden_size
      _lowercase							   =num_hidden_layers
      _lowercase							   =num_attention_heads
      _lowercase							   =intermediate_size
      _lowercase							   =hidden_act
      _lowercase							   =hidden_dropout_prob
      _lowercase							   =attention_probs_dropout_prob
      _lowercase							   =max_position_embeddings
      _lowercase							   =initializer_range
      _lowercase							   =None
      _lowercase							   =vocab_size - 1
      _lowercase							   =vocab_size - 1
      _lowercase							   =vocab_size - 1





     def   __A      (self )  ->     int:
      _lowercase							   =ids_tensor([self.batch_size, self.seq_length]	,  self.vocab_size )

      _lowercase							   =None
      if self.use_input_mask:
       _lowercase							   =random_attention_mask([self.batch_size, self.seq_length] )

      _lowercase							   =GPTaConfig(
          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	,  )

      return (config, input_ids, input_mask)





     def   __A      (self )  ->     List[Any]:
      _lowercase							   =self.prepare_config_and_inputs()
      _lowercase   ,					_lowercase   ,					_lowercase							   =config_and_inputs
      _lowercase							   ={'''input_ids''': input_ids, '''attention_mask''': attention_mask}
      return config, inputs_dict





     def   __A      (self )  ->     Tuple:
      _lowercase   ,					_lowercase   ,					_lowercase							   =self.prepare_config_and_inputs()

      _lowercase							   =floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
      _lowercase							   =ids_tensor([self.batch_size, self.seq_length]	,  vocab_size=2 )

      return (
          config,
          input_ids,
          attention_mask,
          encoder_hidden_states,
          encoder_attention_mask,
      )





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )  ->     Dict:
      _lowercase							   =2_0
      _lowercase							   =model_class_name(UpperCAmelCase )

      _lowercase							   =model.init_cache(input_ids.shape[0]	,  UpperCAmelCase )
      _lowercase							   =jnp.ones((input_ids.shape[0], max_decoder_length)	,  dtype='''i4''' )

      _lowercase							   =jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1 )[None, :]	,  (input_ids.shape[0], input_ids.shape[-1] - 1) )
      _lowercase							   =model(
          input_ids[:, :-1]	,  attention_mask=UpperCAmelCase	,  past_key_values=UpperCAmelCase	,  position_ids=UpperCAmelCase	,  )

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

      _lowercase							   =model(UpperCAmelCase )

      _lowercase							   =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	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )  ->     List[str]:
      _lowercase							   =2_0
      _lowercase							   =model_class_name(UpperCAmelCase )

      _lowercase							   =jnp.concatenate(
          [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1]) )]	,  axis=-1	,  )

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

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

      _lowercase							   =model(UpperCAmelCase	,  attention_mask=UpperCAmelCase )

      _lowercase							   =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__     ( lowerCAmelCase     ,							lowerCAmelCase     ,							unittest.TestCase):
     SCREAMING_SNAKE_CASE__       =    (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
     SCREAMING_SNAKE_CASE__       =    (FlaxGPTaLMHeadModel,) if is_flax_available() else ()





     def   __A      (self )  ->     str:
      _lowercase							   =FlaxGPTaModelTester(self )





     def   __A      (self )  ->     Optional[int]:
      for model_class_name in self.all_model_classes:
       _lowercase   ,					_lowercase   ,					_lowercase							   =self.model_tester.prepare_config_and_inputs()
       self.model_tester.check_use_cache_forward(UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )





     def   __A      (self )  ->     Dict:
      for model_class_name in self.all_model_classes:
       _lowercase   ,					_lowercase   ,					_lowercase							   =self.model_tester.prepare_config_and_inputs()
       self.model_tester.check_use_cache_forward_with_attn_mask(
           UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )





     @slow
     def   __A      (self )  ->     Dict:
      _lowercase							   =GPTaTokenizer.from_pretrained('''gpt2'''	,  pad_token='''</s>'''	,  padding_side='''left''' )
      _lowercase							   =tokenizer(['''Hello this is a long string''', '''Hey''']	,  return_tensors='''np'''	,  padding=UpperCAmelCase	,  truncation=UpperCAmelCase )

      _lowercase							   =FlaxGPTaLMHeadModel.from_pretrained('''gpt2''' )
      _lowercase							   =False
      _lowercase							   =model.config.eos_token_id

      _lowercase							   =jax.jit(model.generate )

      _lowercase							   =jit_generate(inputs['''input_ids''']	,  attention_mask=inputs['''attention_mask'''] ).sequences

      _lowercase							   =tokenizer.batch_decode(UpperCAmelCase	,  skip_special_tokens=UpperCAmelCase )

      _lowercase							   =[
          '''Hello this is a long string of words. I\'m going to start with the first one.\n''',
          '''Hey, I\'m not sure if I\'m going to be able to do''',
      ]

      self.assertListEqual(UpperCAmelCase	,  UpperCAmelCase )





     @is_pt_flax_cross_test
     def   __A      (self )  ->     Optional[Any]:
      _lowercase   ,					_lowercase							   =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
        _lowercase							   =self._prepare_for_class(UpperCAmelCase	,  UpperCAmelCase )
        _lowercase							   ={k: torch.tensor(v.tolist() ) for k, v in prepared_inputs_dict.items()}

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

        _lowercase   ,					_lowercase							   =pt_inputs['''input_ids'''].shape
        _lowercase							   =np.random.randint(0	,  seq_length - 1	,  size=(batch_size,) )
        for batch_idx, start_index in enumerate(UpperCAmelCase ):
         _lowercase							   =0
         _lowercase							   =1
         _lowercase							   =0
         _lowercase							   =1
        _lowercase							   =pt_model_class(UpperCAmelCase ).eval()
        _lowercase							   =model_class(UpperCAmelCase	,  dtype=jnp.floataa )

        _lowercase							   =convert_pytorch_state_dict_to_flax(pt_model.state_dict()	,  UpperCAmelCase )
        _lowercase							   =fx_state

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

        _lowercase							   =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 )
         _lowercase							   =model_class.from_pretrained(UpperCAmelCase	,  from_pt=UpperCAmelCase )

        _lowercase							   =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 )  ->     Optional[Any]:
      _lowercase   ,					_lowercase							   =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
        _lowercase							   =self._prepare_for_class(UpperCAmelCase	,  UpperCAmelCase )
        _lowercase							   ={k: torch.tensor(v.tolist() ) for k, v in prepared_inputs_dict.items()}

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

        _lowercase							   =pt_model_class(UpperCAmelCase ).eval()
        _lowercase							   =model_class(UpperCAmelCase	,  dtype=jnp.floataa )

        _lowercase							   =load_flax_weights_in_pytorch_model(UpperCAmelCase	,  fx_model.params )
        _lowercase   ,					_lowercase							   =pt_inputs['''input_ids'''].shape
        _lowercase							   =np.random.randint(0	,  seq_length - 1	,  size=(batch_size,) )
        for batch_idx, start_index in enumerate(UpperCAmelCase ):
         _lowercase							   =0
         _lowercase							   =1
         _lowercase							   =0
         _lowercase							   =1

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

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

        _lowercase							   =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 )
         _lowercase							   =pt_model_class.from_pretrained(UpperCAmelCase	,  from_flax=UpperCAmelCase )

        with torch.no_grad():
         _lowercase							   =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 )





     @slow
     def   __A      (self )  ->     Union[str, Any]:
      for model_class_name in self.all_model_classes:
       _lowercase							   =model_class_name.from_pretrained('''gpt2'''	,  from_pt=UpperCAmelCase )
       _lowercase							   =model(np.ones((1, 1) ) )
       self.assertIsNotNone(UpperCAmelCase )

