

"""simple docstring"""




from __future__ import annotations

import unittest

from transformers import OpenAIGPTConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
	import tensorflow as tf

	from transformers.models.openai.modeling_tf_openai import (
	    TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TFOpenAIGPTDoubleHeadsModel,
	    TFOpenAIGPTForSequenceClassification,
	    TFOpenAIGPTLMHeadModel,
	    TFOpenAIGPTModel,
	)



class    A__:
						"""simple docstring"""






						def __init__(  self       ,     _lowercase       ,     )      ->				Dict:
									a_		:		Union[str, Any]				    =				parent
									a_		:		Optional[int]				    =				13
									a_		:		Dict				    =				7
									a_		:		Tuple				    =				True
									a_		:		Optional[int]				    =				True
									a_		:		Dict				    =				True
									a_		:		Tuple				    =				True
									a_		:		Optional[int]				    =				True
									a_		:		List[str]				    =				99
									a_		:		List[str]				    =				32
									a_		:		Optional[Any]				    =				2
									a_		:		List[Any]				    =				4
									a_		:		Optional[int]				    =				37
									a_		:		Optional[Any]				    =				"""gelu"""
									a_		:		str				    =				0.1
									a_		:		Optional[Any]				    =				0.1
									a_		:		str				    =				512
									a_		:		List[str]				    =				16
									a_		:		List[str]				    =				2
									a_		:		List[Any]				    =				0.0_2
									a_		:		Union[str, Any]				    =				3
									a_		:		Union[str, Any]				    =				4
									a_		:		Optional[int]				    =				None
									a_		:		Union[str, Any]				    =				self.vocab_size - 1





						def 	UpperCamelCase__			(  self  )      ->				Any:
									a_		:		Tuple				    =				ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size  )

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

									a_		:		List[str]				    =				None
									if self.use_token_type_ids:
												a_		:		Optional[int]				    =				ids_tensor([self.batch_size, self.seq_length]       ,     self.type_vocab_size  )

									a_		:		str				    =				None
									if self.use_mc_token_ids:
												a_		:		str				    =				ids_tensor([self.batch_size, self.num_choices]       ,     self.seq_length  )

									a_		:		int				    =				None
									a_		:		str				    =				None
									a_		:		Optional[int]				    =				None
									if self.use_labels:
												a_		:		Dict				    =				ids_tensor([self.batch_size]       ,     self.type_sequence_label_size  )
												a_		:		str				    =				ids_tensor([self.batch_size, self.seq_length]       ,     self.num_labels  )
												a_		:		Optional[int]				    =				ids_tensor([self.batch_size]       ,     self.num_choices  )

									a_		:		int				    =				OpenAIGPTConfig(
									    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       ,     pad_token_id=self.pad_token_id       ,     )

									a_		:		Dict				    =				ids_tensor([self.num_hidden_layers, self.num_attention_heads]       ,     2  )

									return (
									    config,
									    input_ids,
									    input_mask,
									    head_mask,
									    token_type_ids,
									    mc_token_ids,
									    sequence_labels,
									    token_labels,
									    choice_labels,
									)





						def 	UpperCamelCase__			(  self       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     *_lowercase  )      ->				List[str]:
									a_		:		Union[str, Any]				    =				TFOpenAIGPTModel(config=lowerCAmelCase__  )
									a_		:		List[str]				    =				{"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
									a_		:		Optional[int]				    =				model(lowerCAmelCase__  )

									a_		:		int				    =				[input_ids, input_mask]
									a_		:		Optional[int]				    =				model(lowerCAmelCase__  )

									a_		:		str				    =				model(lowerCAmelCase__  )

									self.parent.assertEqual(result.last_hidden_state.shape       ,     (self.batch_size, self.seq_length, self.hidden_size)  )





						def 	UpperCamelCase__			(  self       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     *_lowercase  )      ->				Union[str, Any]:
									a_		:		Union[str, Any]				    =				TFOpenAIGPTLMHeadModel(config=lowerCAmelCase__  )
									a_		:		Any				    =				{"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
									a_		:		Tuple				    =				model(lowerCAmelCase__  )
									self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.vocab_size)  )





						def 	UpperCamelCase__			(  self       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     *_lowercase  )      ->				str:
									a_		:		Union[str, Any]				    =				TFOpenAIGPTDoubleHeadsModel(config=lowerCAmelCase__  )

									a_		:		Optional[Any]				    =				tf.tile(tf.expand_dims(lowerCAmelCase__       ,     1  )       ,     (1, self.num_choices, 1)  )
									a_		:		Optional[Any]				    =				tf.tile(tf.expand_dims(lowerCAmelCase__       ,     1  )       ,     (1, self.num_choices, 1)  )
									a_		:		str				    =				tf.tile(tf.expand_dims(lowerCAmelCase__       ,     1  )       ,     (1, self.num_choices, 1)  )

									a_		:		str				    =				{
									    """input_ids""": multiple_choice_inputs_ids,
									    """mc_token_ids""": mc_token_ids,
									    """attention_mask""": multiple_choice_input_mask,
									    """token_type_ids""": multiple_choice_token_type_ids,
									}
									a_		:		List[Any]				    =				model(lowerCAmelCase__  )
									self.parent.assertEqual(
									    result.logits.shape       ,     (self.batch_size, self.num_choices, self.seq_length, self.vocab_size)  )
									self.parent.assertEqual(result.mc_logits.shape       ,     (self.batch_size, self.num_choices)  )





						def 	UpperCamelCase__			(  self       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     *_lowercase  )      ->				Union[str, Any]:
									a_		:		Tuple				    =				self.num_labels
									a_		:		Optional[Any]				    =				ids_tensor([self.batch_size]       ,     self.type_sequence_label_size  )
									a_		:		Optional[Any]				    =				{
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									    """labels""": sequence_labels,
									}
									a_		:		Union[str, Any]				    =				TFOpenAIGPTForSequenceClassification(lowerCAmelCase__  )
									a_		:		Any				    =				model(lowerCAmelCase__  )
									self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.num_labels)  )





						def 	UpperCamelCase__			(  self  )      ->				Dict:
									a_		:		Tuple				    =				self.prepare_config_and_inputs()

									(
									    (
									    a_ 
									)    ,			(
									    a_ 
									)    ,			(
									    a_ 
									)    ,			(
									    a_ 
									)    ,			(
									    a_ 
									)    ,			(
									    a_ 
									)    ,			(
									    a_ 
									)    ,			(
									    a_ 
									)    ,			(
									    a_ 
									)    ,			 
									)		:		int				    =				config_and_inputs

									a_		:		Tuple				    =				{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
									return config, inputs_dict







@require_tf
class    A__(lowercase__,       lowercase__,       unittest.TestCase    ):
						"""simple docstring"""






						_A  :							Optional[int]   			=	(
						    (TFOpenAIGPTModel, TFOpenAIGPTLMHeadModel, TFOpenAIGPTDoubleHeadsModel, TFOpenAIGPTForSequenceClassification)
						    if is_tf_available()
						    else ()
						)
						_A  :							Dict   			=	(
						    (TFOpenAIGPTLMHeadModel,) if is_tf_available() else ()
						)  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
						_A  :							Union[str, Any]   			=	(
						    {
						        '''feature-extraction''': TFOpenAIGPTModel,
						        '''text-classification''': TFOpenAIGPTForSequenceClassification,
						        '''text-generation''': TFOpenAIGPTLMHeadModel,
						        '''zero-shot''': TFOpenAIGPTForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						_A  :							int   			=	False
						_A  :							Optional[Any]   			=	False





						def 	UpperCamelCase__			(  self       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase  )      ->				Optional[Any]:
									if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
												# Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
												# `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
												# tiny config could not be created.
												return True

									return False





						def 	UpperCamelCase__			(  self  )      ->				Union[str, Any]:
									a_		:		int				    =				TFOpenAIGPTModelTester(self  )
									a_		:		str				    =				ConfigTester(self       ,     config_class=lowerCAmelCase__       ,     n_embd=37  )





						def 	UpperCamelCase__			(  self  )      ->				Dict:
									self.config_tester.run_common_tests()





						def 	UpperCamelCase__			(  self  )      ->				Tuple:
									a_		:		Optional[Any]				    =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_model(*lowerCAmelCase__  )





						def 	UpperCamelCase__			(  self  )      ->				Tuple:
									a_		:		int				    =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_lm_head(*lowerCAmelCase__  )





						def 	UpperCamelCase__			(  self  )      ->				int:
									a_		:		List[Any]				    =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_double_head(*lowerCAmelCase__  )





						def 	UpperCamelCase__			(  self  )      ->				Dict:
									a_		:		Tuple				    =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*lowerCAmelCase__  )





						@slow
						def 	UpperCamelCase__			(  self  )      ->				Union[str, Any]:
									for model_name in TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												a_		:		Optional[int]				    =				TFOpenAIGPTModel.from_pretrained(lowerCAmelCase__  )
												self.assertIsNotNone(lowerCAmelCase__  )







@require_tf
class    A__(unittest.TestCase    ):
						"""simple docstring"""






						@slow
						def 	UpperCamelCase__			(  self  )      ->				List[str]:
									a_		:		List[Any]				    =				TFOpenAIGPTLMHeadModel.from_pretrained("""openai-gpt"""  )
									a_		:		Dict				    =				tf.convert_to_tensor([[481, 4_735, 544]]       ,     dtype=tf.intaa  )  # the president is
									a_		:		Tuple				    =				[
									    481,
									    4_735,
									    544,
									    246,
									    963,
									    870,
									    762,
									    239,
									    244,
									    40_477,
									    244,
									    249,
									    719,
									    881,
									    487,
									    544,
									    240,
									    244,
									    603,
									    481,
									]  # the president is a very good man. " \n " i\'m sure he is, " said the

									a_		:		Dict				    =				model.generate(lowerCAmelCase__       ,     do_sample=lowerCAmelCase__  )
									self.assertListEqual(output_ids[0].numpy().tolist()       ,     lowerCAmelCase__  )





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    A__:
						"""simple docstring"""






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





						def 	UpperCamelCase__			(  self  )      ->				Optional[int]:
									a_		:		Dict				    =				ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size  )

									a_		:		List[str]				    =				None
									if self.use_input_mask:
												a_		:		Dict				    =				random_attention_mask([self.batch_size, self.seq_length]  )

									a_		:		Optional[int]				    =				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=_lowercase       ,     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 	UpperCamelCase__			(  self  )      ->				int:
									a_		:		Any				    =				self.prepare_config_and_inputs()
									a_    ,			a_    ,			a_		:		List[str]				    =				config_and_inputs
									a_		:		Optional[Any]				    =				{"""input_ids""": input_ids, """attention_mask""": attention_mask}
									return config, inputs_dict





						def 	UpperCamelCase__			(  self  )      ->				Tuple:
									a_    ,			a_    ,			a_		:		Dict				    =				self.prepare_config_and_inputs()

									a_		:		List[str]				    =				floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
									a_		:		Optional[int]				    =				ids_tensor([self.batch_size, self.seq_length]       ,     vocab_size=2  )

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





						def 	UpperCamelCase__			(  self       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase  )      ->				Optional[Any]:
									a_		:		Dict				    =				20
									a_		:		Tuple				    =				model_class_name(_lowercase  )

									a_		:		Tuple				    =				model.init_cache(input_ids.shape[0]       ,     _lowercase  )
									a_		:		Any				    =				jnp.ones((input_ids.shape[0], max_decoder_length)       ,     dtype="""i4"""  )

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

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

									a_		:		Dict				    =				model(_lowercase  )

									a_		:		Tuple				    =				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 	UpperCamelCase__			(  self       ,     _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase  )      ->				int:
									a_		:		Any				    =				20
									a_		:		Optional[int]				    =				model_class_name(_lowercase  )

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

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

									a_		:		Optional[Any]				    =				model(
									    input_ids[:, :-1]       ,     attention_mask=_lowercase       ,     past_key_values=_lowercase       ,     position_ids=_lowercase       ,     )
									a_		:		List[str]				    =				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,     dtype="""i4"""  )
									a_		:		Dict				    =				model(
									    input_ids[:, -1:]       ,     past_key_values=outputs_cache.past_key_values       ,     attention_mask=_lowercase       ,     position_ids=_lowercase       ,     )

									a_		:		str				    =				model(_lowercase       ,     attention_mask=_lowercase  )

									a_		:		List[str]				    =				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    A__(a_,       a_,       unittest.TestCase    ):
						"""simple docstring"""






						_A  :							List[str]   			=	(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
						_A  :							Dict   			=	(FlaxGPTaLMHeadModel,) if is_flax_available() else ()





						def 	UpperCamelCase__			(  self  )      ->				Optional[Any]:
									a_		:		List[str]				    =				FlaxGPTaModelTester(self  )





						def 	UpperCamelCase__			(  self  )      ->				int:
									for model_class_name in self.all_model_classes:
												a_    ,			a_    ,			a_		:		Any				    =				self.model_tester.prepare_config_and_inputs()
												self.model_tester.check_use_cache_forward(_lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase  )





						def 	UpperCamelCase__			(  self  )      ->				int:
									for model_class_name in self.all_model_classes:
												a_    ,			a_    ,			a_		:		Optional[Any]				    =				self.model_tester.prepare_config_and_inputs()
												self.model_tester.check_use_cache_forward_with_attn_mask(
												    _lowercase       ,     _lowercase       ,     _lowercase       ,     _lowercase  )





						@slow
						def 	UpperCamelCase__			(  self  )      ->				Tuple:
									a_		:		str				    =				GPTaTokenizer.from_pretrained("""gpt2"""       ,     pad_token="""</s>"""       ,     padding_side="""left"""  )
									a_		:		Optional[int]				    =				tokenizer(["""Hello this is a long string""", """Hey"""]       ,     return_tensors="""np"""       ,     padding=_lowercase       ,     truncation=_lowercase  )

									a_		:		Any				    =				FlaxGPTaLMHeadModel.from_pretrained("""gpt2"""  )
									a_		:		List[Any]				    =				False
									a_		:		List[Any]				    =				model.config.eos_token_id

									a_		:		Tuple				    =				jax.jit(model.generate  )

									a_		:		Union[str, Any]				    =				jit_generate(inputs["""input_ids"""]       ,     attention_mask=inputs["""attention_mask"""]  ).sequences

									a_		:		Dict				    =				tokenizer.batch_decode(_lowercase       ,     skip_special_tokens=_lowercase  )

									a_		:		Tuple				    =				[
									    """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(_lowercase       ,     _lowercase  )





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

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

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

															a_		:		Optional[int]				    =				convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,     _lowercase  )
															a_		:		Union[str, Any]				    =				fx_state

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

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

															with tempfile.TemporaryDirectory() as tmpdirname:
																		pt_model.save_pretrained(_lowercase  )
																		a_		:		Dict				    =				model_class.from_pretrained(_lowercase       ,     from_pt=_lowercase  )

															a_		:		int				    =				fx_model_loaded(**_lowercase  ).to_tuple()
															self.assertEqual(
															    len(_lowercase  )       ,     len(_lowercase  )       ,     """Output lengths differ between Flax and PyTorch"""  )
															for fx_output_loaded, pt_output in zip(_lowercase       ,     _lowercase  ):
																		self.assert_almost_equals(fx_output_loaded[:, -1]       ,     pt_output[:, -1].numpy()       ,     4e-2  )





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

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

															a_		:		Tuple				    =				pt_model_class(_lowercase  ).eval()
															a_		:		Dict				    =				model_class(_lowercase       ,     dtype=jnp.floataa  )

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

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

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

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

															with tempfile.TemporaryDirectory() as tmpdirname:
																		fx_model.save_pretrained(_lowercase  )
																		a_		:		Dict				    =				pt_model_class.from_pretrained(_lowercase       ,     from_flax=_lowercase  )

															with torch.no_grad():
																		a_		:		List[str]				    =				pt_model_loaded(**_lowercase  ).to_tuple()

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





						@slow
						def 	UpperCamelCase__			(  self  )      ->				Any:
									for model_class_name in self.all_model_classes:
												a_		:		List[str]				    =				model_class_name.from_pretrained("""gpt2"""       ,     from_pt=_lowercase  )
												a_		:		List[str]				    =				model(np.ones((1, 1)  )  )
												self.assertIsNotNone(_lowercase  )


