



from __future__ import annotations

import inspect
import unittest
from typing import List, Tuple

from transformers import ConvNextConfig
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 TFConvNextForImageClassification, TFConvNextModel


if is_vision_available():
	from PIL import Image

	from transformers import ConvNextImageProcessor


class  _lowercase				:
				'''simple docstring'''







				def __init__(	self       ,	snake_case__       ,	snake_case__=13       ,	snake_case__=32       ,	snake_case__=3       ,	snake_case__=4       ,	snake_case__=[10, 20, 30, 40]       ,	snake_case__=[2, 2, 3, 2]       ,	snake_case__=True       ,	snake_case__=True       ,	snake_case__=37       ,	snake_case__="gelu"       ,	snake_case__=10       ,	snake_case__=0.02       ,	snake_case__=3       ,	snake_case__=None       ,	):

									'''simple docstring'''


									UpperCamelCase_											=				parent
									UpperCamelCase_											=				batch_size
									UpperCamelCase_											=				image_size
									UpperCamelCase_											=				num_channels
									UpperCamelCase_											=				num_stages
									UpperCamelCase_											=				hidden_sizes
									UpperCamelCase_											=				depths
									UpperCamelCase_											=				is_training
									UpperCamelCase_											=				use_labels
									UpperCamelCase_											=				intermediate_size
									UpperCamelCase_											=				hidden_act
									UpperCamelCase_											=				type_sequence_label_size
									UpperCamelCase_											=				initializer_range
									UpperCamelCase_											=				scope

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

									UpperCamelCase_											=				None
									if self.use_labels:
														UpperCamelCase_											=				ids_tensor([self.batch_size]       ,	self.type_sequence_label_size      )

									UpperCamelCase_											=				self.get_config()

									return config, pixel_values, labels

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									return ConvNextConfig(
									    num_channels=self.num_channels       ,	hidden_sizes=self.hidden_sizes       ,	depths=self.depths       ,	num_stages=self.num_stages       ,	hidden_act=self.hidden_act       ,	is_decoder=snake_case__       ,	initializer_range=self.initializer_range       ,	)

				def  _lowerCamelCase   (	self       ,	snake_case__       ,	snake_case__       ,	snake_case__      ):

									'''simple docstring'''


									UpperCamelCase_											=				TFConvNextModel(config=snake_case__      )
									UpperCamelCase_											=				model(snake_case__       ,	training=snake_case__      )
									# 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  _lowerCamelCase   (	self       ,	snake_case__       ,	snake_case__       ,	snake_case__      ):

									'''simple docstring'''


									UpperCamelCase_											=				self.type_sequence_label_size
									UpperCamelCase_											=				TFConvNextForImageClassification(snake_case__      )
									UpperCamelCase_											=				model(snake_case__       ,	labels=snake_case__       ,	training=snake_case__      )
									self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.type_sequence_label_size)      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				self.prepare_config_and_inputs()
									UpperCamelCase_					,      UpperCamelCase_					,      UpperCamelCase_											=				config_and_inputs
									UpperCamelCase_											=				{"pixel_values": pixel_values}
									return config, inputs_dict




@require_tf
class  _lowercase				(a_ ,				a_ ,				unittest.TestCase ):
				'''simple docstring'''







				lowercase__				  =  (TFConvNextModel, TFConvNextForImageClassification) if is_tf_available() else ()
				lowercase__				  =  (
				    {"""feature-extraction""": TFConvNextModel, """image-classification""": TFConvNextForImageClassification}
				    if is_tf_available()
				    else {}
				)

				lowercase__				  =  False
				lowercase__				  =  False
				lowercase__				  =  False
				lowercase__				  =  False
				lowercase__				  =  False

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				TFConvNextModelTester(self      )
									UpperCamelCase_											=				ConfigTester(
									    self       ,	config_class=snake_case__       ,	has_text_modality=snake_case__       ,	hidden_size=37       ,	)

				@unittest.skip(reason="ConvNext does not use inputs_embeds"      )
				def  _lowerCamelCase   (	self      ):

									'''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  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									super().test_keras_fit()

				@unittest.skip(reason="ConvNext does not support input and output embeddings"      )
				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									pass

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_					,      UpperCamelCase_											=				self.model_tester.prepare_config_and_inputs_for_common()

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

														UpperCamelCase_											=				["pixel_values"]
														self.assertListEqual(arg_names[:1]       ,	snake_case__      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*snake_case__      )

				@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."       ,	)
				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									super().test_dataset_conversion()

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									def check_hidden_states_output(snake_case__       ,	snake_case__       ,	snake_case__      ):
														UpperCamelCase_											=				model_class(snake_case__      )

														UpperCamelCase_											=				model(**self._prepare_for_class(snake_case__       ,	snake_case__      )      )
														UpperCamelCase_											=				outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

														UpperCamelCase_											=				self.model_tester.num_stages
														self.assertEqual(len(snake_case__      )       ,	expected_num_stages + 1      )

														# ConvNext'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 // 4, self.model_tester.image_size // 4]       ,	)

									UpperCamelCase_					,      UpperCamelCase_											=				self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
														UpperCamelCase_											=				True
														check_hidden_states_output(snake_case__       ,	snake_case__       ,	snake_case__      )

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

														check_hidden_states_output(snake_case__       ,	snake_case__       ,	snake_case__      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_					,      UpperCamelCase_											=				self.model_tester.prepare_config_and_inputs_for_common()

									def check_equivalence(snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__={}      ):
														UpperCamelCase_											=				model(snake_case__       ,	return_dict=snake_case__       ,	**snake_case__      )
														UpperCamelCase_											=				model(snake_case__       ,	return_dict=snake_case__       ,	**snake_case__      ).to_tuple()

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

																			recursive_check(snake_case__       ,	snake_case__      )

									for model_class in self.all_model_classes:
														UpperCamelCase_											=				model_class(snake_case__      )

														UpperCamelCase_											=				self._prepare_for_class(snake_case__       ,	snake_case__      )
														UpperCamelCase_											=				self._prepare_for_class(snake_case__       ,	snake_case__      )
														check_equivalence(snake_case__       ,	snake_case__       ,	snake_case__      )

														UpperCamelCase_											=				self._prepare_for_class(snake_case__       ,	snake_case__       ,	return_labels=snake_case__      )
														UpperCamelCase_											=				self._prepare_for_class(snake_case__       ,	snake_case__       ,	return_labels=snake_case__      )
														check_equivalence(snake_case__       ,	snake_case__       ,	snake_case__      )

														UpperCamelCase_											=				self._prepare_for_class(snake_case__       ,	snake_case__      )
														UpperCamelCase_											=				self._prepare_for_class(snake_case__       ,	snake_case__      )
														check_equivalence(snake_case__       ,	snake_case__       ,	snake_case__       ,	{"output_hidden_states": True}      )

														UpperCamelCase_											=				self._prepare_for_class(snake_case__       ,	snake_case__       ,	return_labels=snake_case__      )
														UpperCamelCase_											=				self._prepare_for_class(snake_case__       ,	snake_case__       ,	return_labels=snake_case__      )
														check_equivalence(snake_case__       ,	snake_case__       ,	snake_case__       ,	{"output_hidden_states": True}      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*snake_case__      )

				@slow
				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				TFConvNextModel.from_pretrained("facebook/convnext-tiny-224"      )
									self.assertIsNotNone(snake_case__      )




def 				_lowerCAmelCase      ():
					UpperCamelCase_											=				Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
					return image






@require_tf
@require_vision
class  _lowercase				(unittest.TestCase ):
				'''simple docstring'''







				@cached_property
				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									return ConvNextImageProcessor.from_pretrained("facebook/convnext-tiny-224"      ) if is_vision_available() else None

				@slow
				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				TFConvNextForImageClassification.from_pretrained("facebook/convnext-tiny-224"      )

									UpperCamelCase_											=				self.default_image_processor
									UpperCamelCase_											=				prepare_img()
									UpperCamelCase_											=				image_processor(images=snake_case__       ,	return_tensors="tf"      )

									# forward pass
									UpperCamelCase_											=				model(**snake_case__      )

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

									UpperCamelCase_											=				tf.constant([-0.0_260, -0.4_739, 0.1_911]      )

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





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  _lowercase				:
				'''simple docstring'''







				def __init__(	self       ,	snake_case__       ,	snake_case__=14       ,	snake_case__=7       ,	snake_case__=True       ,	snake_case__=True       ,	snake_case__=False       ,	snake_case__=True       ,	snake_case__=99       ,	snake_case__=32       ,	snake_case__=5       ,	snake_case__=4       ,	snake_case__=37       ,	snake_case__="gelu"       ,	snake_case__=0.1       ,	snake_case__=0.1       ,	snake_case__=512       ,	snake_case__=0.02       ,	):

									'''simple docstring'''


									UpperCamelCase_											=				parent
									UpperCamelCase_											=				batch_size
									UpperCamelCase_											=				seq_length
									UpperCamelCase_											=				is_training
									UpperCamelCase_											=				use_input_mask
									UpperCamelCase_											=				use_token_type_ids
									UpperCamelCase_											=				use_labels
									UpperCamelCase_											=				vocab_size
									UpperCamelCase_											=				hidden_size
									UpperCamelCase_											=				num_hidden_layers
									UpperCamelCase_											=				num_attention_heads
									UpperCamelCase_											=				intermediate_size
									UpperCamelCase_											=				hidden_act
									UpperCamelCase_											=				hidden_dropout_prob
									UpperCamelCase_											=				attention_probs_dropout_prob
									UpperCamelCase_											=				max_position_embeddings
									UpperCamelCase_											=				initializer_range
									UpperCamelCase_											=				None
									UpperCamelCase_											=				vocab_size - 1
									UpperCamelCase_											=				vocab_size - 1
									UpperCamelCase_											=				vocab_size - 1

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


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

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

									UpperCamelCase_											=				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=snake_case__       ,	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  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				self.prepare_config_and_inputs()
									UpperCamelCase_					,      UpperCamelCase_					,      UpperCamelCase_											=				config_and_inputs
									UpperCamelCase_											=				{"input_ids": input_ids, "attention_mask": attention_mask}
									return config, inputs_dict

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_					,      UpperCamelCase_					,      UpperCamelCase_											=				self.prepare_config_and_inputs()

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

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

				def  _lowerCamelCase   (	self       ,	snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__      ):

									'''simple docstring'''


									UpperCamelCase_											=				20
									UpperCamelCase_											=				model_class_name(snake_case__      )

									UpperCamelCase_											=				model.init_cache(input_ids.shape[0]       ,	snake_case__      )
									UpperCamelCase_											=				jnp.ones((input_ids.shape[0], max_decoder_length)       ,	dtype="i4"      )

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

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

									UpperCamelCase_											=				model(snake_case__      )

									UpperCamelCase_											=				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  _lowerCamelCase   (	self       ,	snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__      ):

									'''simple docstring'''


									UpperCamelCase_											=				20
									UpperCamelCase_											=				model_class_name(snake_case__      )

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

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

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

									UpperCamelCase_											=				model(snake_case__       ,	attention_mask=snake_case__      )

									UpperCamelCase_											=				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  _lowercase				(a_ ,				a_ ,				unittest.TestCase ):
				'''simple docstring'''







				lowercase__				  =  (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
				lowercase__				  =  (FlaxGPTaLMHeadModel,) if is_flax_available() else ()

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				FlaxGPTaModelTester(self      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									for model_class_name in self.all_model_classes:
														UpperCamelCase_					,      UpperCamelCase_					,      UpperCamelCase_											=				self.model_tester.prepare_config_and_inputs()
														self.model_tester.check_use_cache_forward(snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__      )

				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									for model_class_name in self.all_model_classes:
														UpperCamelCase_					,      UpperCamelCase_					,      UpperCamelCase_											=				self.model_tester.prepare_config_and_inputs()
														self.model_tester.check_use_cache_forward_with_attn_mask(
														    snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__      )

				@slow
				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_											=				GPTaTokenizer.from_pretrained("gpt2"       ,	pad_token="</s>"       ,	padding_side="left"      )
									UpperCamelCase_											=				tokenizer(["Hello this is a long string", "Hey"]       ,	return_tensors="np"       ,	padding=snake_case__       ,	truncation=snake_case__      )

									UpperCamelCase_											=				FlaxGPTaLMHeadModel.from_pretrained("gpt2"      )
									UpperCamelCase_											=				False
									UpperCamelCase_											=				model.config.eos_token_id

									UpperCamelCase_											=				jax.jit(model.generate      )

									UpperCamelCase_											=				jit_generate(inputs["input_ids"]       ,	attention_mask=inputs["attention_mask"]      ).sequences

									UpperCamelCase_											=				tokenizer.batch_decode(snake_case__       ,	skip_special_tokens=snake_case__      )

									UpperCamelCase_											=				[
									    "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(snake_case__       ,	snake_case__      )

				@is_pt_flax_cross_test
				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_					,      UpperCamelCase_											=				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
																			UpperCamelCase_											=				self._prepare_for_class(snake_case__       ,	snake_case__      )
																			UpperCamelCase_											=				{k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

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

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

																			UpperCamelCase_											=				convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,	snake_case__      )
																			UpperCamelCase_											=				fx_state

																			with torch.no_grad():
																								UpperCamelCase_											=				pt_model(**snake_case__      ).to_tuple()

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

																			with tempfile.TemporaryDirectory() as tmpdirname:
																								pt_model.save_pretrained(snake_case__      )
																								UpperCamelCase_											=				model_class.from_pretrained(snake_case__       ,	from_pt=snake_case__      )

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

				@is_pt_flax_cross_test
				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									UpperCamelCase_					,      UpperCamelCase_											=				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
																			UpperCamelCase_											=				self._prepare_for_class(snake_case__       ,	snake_case__      )
																			UpperCamelCase_											=				{k: torch.tensor(v.tolist()      ) for k, v in prepared_inputs_dict.items()}

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

																			UpperCamelCase_											=				pt_model_class(snake_case__      ).eval()
																			UpperCamelCase_											=				model_class(snake_case__       ,	dtype=jnp.floataa      )

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

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

																			with torch.no_grad():
																								UpperCamelCase_											=				pt_model(**snake_case__      ).to_tuple()

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

																			with tempfile.TemporaryDirectory() as tmpdirname:
																								fx_model.save_pretrained(snake_case__      )
																								UpperCamelCase_											=				pt_model_class.from_pretrained(snake_case__       ,	from_flax=snake_case__      )

																			with torch.no_grad():
																								UpperCamelCase_											=				pt_model_loaded(**snake_case__      ).to_tuple()

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

				@slow
				def  _lowerCamelCase   (	self      ):

									'''simple docstring'''


									for model_class_name in self.all_model_classes:
														UpperCamelCase_											=				model_class_name.from_pretrained("gpt2"       ,	from_pt=snake_case__      )
														UpperCamelCase_											=				model(np.ones((1, 1)      )      )
														self.assertIsNotNone(snake_case__      )
