




import collections
import inspect
import unittest

from transformers import FocalNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch
					from torch import nn

					from transformers import (
					    FocalNetBackbone,
					    FocalNetForImageClassification,
					    FocalNetForMaskedImageModeling,
					    FocalNetModel,
					)
					from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
					from PIL import Image

					from transformers import AutoImageProcessor
class       a__   :






						def __init__(       self		,      _A		,      _A=1_3		,      _A=3_2		,      _A=2		,      _A=3		,      _A=1_6		,      _A=[3_2, 6_4, 1_2_8]		,      _A=[1, 2, 1]		,      _A=[2, 2, 4]		,      _A=2		,      _A=2.0		,      _A=True		,      _A=0.0		,      _A=0.0		,      _A=0.1		,      _A="gelu"		,      _A=False		,      _A=True		,      _A=0.02		,      _A=1E-5		,      _A=True		,      _A=None		,      _A=True		,      _A=1_0		,      _A=8		,      _A=["stage1", "stage2"]		,      _A=[1, 2]		,      ):




											"""simple docstring"""



											__lowerCAmelCase           =	parent
											__lowerCAmelCase           =	batch_size
											__lowerCAmelCase           =	image_size
											__lowerCAmelCase           =	patch_size
											__lowerCAmelCase           =	num_channels
											__lowerCAmelCase           =	embed_dim
											__lowerCAmelCase           =	hidden_sizes
											__lowerCAmelCase           =	depths
											__lowerCAmelCase           =	num_heads
											__lowerCAmelCase           =	window_size
											__lowerCAmelCase           =	mlp_ratio
											__lowerCAmelCase           =	qkv_bias
											__lowerCAmelCase           =	hidden_dropout_prob
											__lowerCAmelCase           =	attention_probs_dropout_prob
											__lowerCAmelCase           =	drop_path_rate
											__lowerCAmelCase           =	hidden_act
											__lowerCAmelCase           =	use_absolute_embeddings
											__lowerCAmelCase           =	patch_norm
											__lowerCAmelCase           =	layer_norm_eps
											__lowerCAmelCase           =	initializer_range
											__lowerCAmelCase           =	is_training
											__lowerCAmelCase           =	scope
											__lowerCAmelCase           =	use_labels
											__lowerCAmelCase           =	type_sequence_label_size
											__lowerCAmelCase           =	encoder_stride
											__lowerCAmelCase           =	out_features
											__lowerCAmelCase           =	out_indices






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



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

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

											__lowerCAmelCase           =	self.get_config()

											return config, pixel_values, labels






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return FocalNetConfig(
											    image_size=self.image_size		,      patch_size=self.patch_size		,      num_channels=self.num_channels		,      embed_dim=self.embed_dim		,      hidden_sizes=self.hidden_sizes		,      depths=self.depths		,      num_heads=self.num_heads		,      window_size=self.window_size		,      mlp_ratio=self.mlp_ratio		,      qkv_bias=self.qkv_bias		,      hidden_dropout_prob=self.hidden_dropout_prob		,      attention_probs_dropout_prob=self.attention_probs_dropout_prob		,      drop_path_rate=self.drop_path_rate		,      hidden_act=self.hidden_act		,      use_absolute_embeddings=self.use_absolute_embeddings		,      path_norm=self.patch_norm		,      layer_norm_eps=self.layer_norm_eps		,      initializer_range=self.initializer_range		,      encoder_stride=self.encoder_stride		,      out_features=self.out_features		,      out_indices=self.out_indices		,      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	FocalNetModel(config=_A      )
											model.to(_A      )
											model.eval()
											__lowerCAmelCase           =	model(_A      )

											__lowerCAmelCase           =	((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths      ) - 1))
											__lowerCAmelCase           =	int(config.embed_dim * 2 ** (len(config.depths      ) - 1)      )

											self.parent.assertEqual(result.last_hidden_state.shape		,      (self.batch_size, expected_seq_len, expected_dim)      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	FocalNetBackbone(config=_A      )
											model.to(_A      )
											model.eval()
											__lowerCAmelCase           =	model(_A      )

											# verify feature maps
											self.parent.assertEqual(len(result.feature_maps      )		,      len(config.out_features      )      )
											self.parent.assertListEqual(list(result.feature_maps[0].shape      )		,      [self.batch_size, self.image_size, 8, 8]      )

											# verify channels
											self.parent.assertEqual(len(model.channels      )		,      len(config.out_features      )      )
											self.parent.assertListEqual(model.channels		,      config.hidden_sizes[:-1]      )

											# verify backbone works with out_features=None
											__lowerCAmelCase           =	None
											__lowerCAmelCase           =	FocalNetBackbone(config=_A      )
											model.to(_A      )
											model.eval()
											__lowerCAmelCase           =	model(_A      )

											# verify feature maps
											self.parent.assertEqual(len(result.feature_maps      )		,      1      )
											self.parent.assertListEqual(list(result.feature_maps[0].shape      )		,      [self.batch_size, self.image_size * 2, 4, 4]      )

											# verify channels
											self.parent.assertEqual(len(model.channels      )		,      1      )
											self.parent.assertListEqual(model.channels		,      [config.hidden_sizes[-1]]      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	FocalNetForMaskedImageModeling(config=_A      )
											model.to(_A      )
											model.eval()
											__lowerCAmelCase           =	model(_A      )
											self.parent.assertEqual(
											    result.reconstruction.shape		,      (self.batch_size, self.num_channels, self.image_size, self.image_size)      )

											# test greyscale images
											__lowerCAmelCase           =	1
											__lowerCAmelCase           =	FocalNetForMaskedImageModeling(_A      )
											model.to(_A      )
											model.eval()

											__lowerCAmelCase           =	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
											__lowerCAmelCase           =	model(_A      )
											self.parent.assertEqual(result.reconstruction.shape		,      (self.batch_size, 1, self.image_size, self.image_size)      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.type_sequence_label_size
											__lowerCAmelCase           =	FocalNetForImageClassification(_A      )
											model.to(_A      )
											model.eval()
											__lowerCAmelCase           =	model(_A		,      labels=_A      )
											self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.type_sequence_label_size)      )

											# test greyscale images
											__lowerCAmelCase           =	1
											__lowerCAmelCase           =	FocalNetForImageClassification(_A      )
											model.to(_A      )
											model.eval()

											__lowerCAmelCase           =	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
											__lowerCAmelCase           =	model(_A      )
											self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.type_sequence_label_size)      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.prepare_config_and_inputs()

											__lowerCAmelCase       ,							__lowerCAmelCase       ,							__lowerCAmelCase           =	config_and_inputs
											__lowerCAmelCase           =	{"pixel_values": pixel_values}
											return config, inputs_dict






@require_torch
class       a__   ( snake_case__			,    snake_case__			,    unittest.TestCase				):
						_a							:   Union[str, Any] =   (
						    (
						        FocalNetModel,
						        FocalNetForImageClassification,
						        FocalNetForMaskedImageModeling,
						        FocalNetBackbone,
						    )
						    if is_torch_available()
						    else ()
						)
						_a							:   str =   (
						    {"""feature-extraction""": FocalNetModel, """image-classification""": FocalNetForImageClassification}
						    if is_torch_available()
						    else {}
						)
						_a							:   Tuple =   False

						_a							:   str =   False
						_a							:   int =   False
						_a							:   Optional[Any] =   False
						_a							:   int =   False






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	FocalNetModelTester(self      )
											__lowerCAmelCase           =	ConfigTester(self		,      config_class=_A		,      embed_dim=3_7		,      has_text_modality=_A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											self.create_and_test_config_common_properties()
											self.config_tester.create_and_test_config_to_json_string()
											self.config_tester.create_and_test_config_to_json_file()
											self.config_tester.create_and_test_config_from_and_save_pretrained()
											self.config_tester.create_and_test_config_with_num_labels()
											self.config_tester.check_config_can_be_init_without_params()
											self.config_tester.check_config_arguments_init()






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_backbone(*_A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_image_modeling(*_A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*_A      )






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




											"""simple docstring"""



											pass






						@unittest.skip(reason="FocalNet does not use feedforward chunking"      )
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											pass






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes[:-1]:
																__lowerCAmelCase           =	model_class(_A      )
																self.assertIsInstance(model.get_input_embeddings()		,      (nn.Module)      )
																__lowerCAmelCase           =	model.get_output_embeddings()
																self.assertTrue(x is None or isinstance(_A		,      nn.Linear      )      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes[:-1]:
																__lowerCAmelCase           =	model_class(_A      )
																__lowerCAmelCase           =	inspect.signature(model.forward      )
																# signature.parameters is an OrderedDict => so arg_names order is deterministic
																__lowerCAmelCase           =	[*signature.parameters.keys()]

																__lowerCAmelCase           =	["pixel_values"]
																self.assertListEqual(arg_names[:1]		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	model_class(_A      )
											model.to(_A      )
											model.eval()

											with torch.no_grad():
																__lowerCAmelCase           =	model(**self._prepare_for_class(_A		,      _A      )      )

											__lowerCAmelCase           =	outputs.hidden_states

											__lowerCAmelCase           =	getattr(
											    self.model_tester		,      "expected_num_hidden_layers"		,      len(self.model_tester.depths      ) + 1      )
											self.assertEqual(len(_A      )		,      _A      )

											# FocalNet has a different seq_length
											__lowerCAmelCase           =	(
											    config.patch_size
											    if isinstance(config.patch_size		,      collections.abc.Iterable      )
											    else (config.patch_size, config.patch_size)
											)

											__lowerCAmelCase           =	(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

											self.assertListEqual(
											    list(hidden_states[0].shape[-2:]      )		,      [num_patches, self.model_tester.embed_dim]		,      )

											__lowerCAmelCase           =	outputs.reshaped_hidden_states
											self.assertEqual(len(_A      )		,      _A      )

											__lowerCAmelCase       ,							__lowerCAmelCase       ,							__lowerCAmelCase       ,							__lowerCAmelCase           =	reshaped_hidden_states[0].shape
											__lowerCAmelCase           =	(
											    reshaped_hidden_states[0].view(_A		,      _A		,      height * width      ).permute(0		,      2		,      1      )
											)
											self.assertListEqual(
											    list(reshaped_hidden_states.shape[-2:]      )		,      [num_patches, self.model_tester.embed_dim]		,      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs_for_common()

											__lowerCAmelCase           =	(
											    self.model_tester.image_size
											    if isinstance(self.model_tester.image_size		,      collections.abc.Iterable      )
											    else (self.model_tester.image_size, self.model_tester.image_size)
											)

											for model_class in self.all_model_classes[:-1]:
																__lowerCAmelCase           =	True
																self.check_hidden_states_output(_A		,      _A		,      _A		,      _A      )

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

																self.check_hidden_states_output(_A		,      _A		,      _A		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs_for_common()
											__lowerCAmelCase           =	3

											__lowerCAmelCase           =	(
											    self.model_tester.image_size
											    if isinstance(self.model_tester.image_size		,      collections.abc.Iterable      )
											    else (self.model_tester.image_size, self.model_tester.image_size)
											)
											__lowerCAmelCase           =	(
											    config.patch_size
											    if isinstance(config.patch_size		,      collections.abc.Iterable      )
											    else (config.patch_size, config.patch_size)
											)

											__lowerCAmelCase           =	image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
											__lowerCAmelCase           =	image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

											for model_class in self.all_model_classes[:-1]:
																__lowerCAmelCase           =	True
																self.check_hidden_states_output(_A		,      _A		,      _A		,      (padded_height, padded_width)      )

																# check that output_hidden_states also work using config
																del inputs_dict["output_hidden_states"]
																__lowerCAmelCase           =	True
																self.check_hidden_states_output(_A		,      _A		,      _A		,      (padded_height, padded_width)      )






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__lowerCAmelCase           =	FocalNetModel.from_pretrained(_A      )
																self.assertIsNotNone(_A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs_for_common()

											__lowerCAmelCase           =	_config_zero_init(_A      )
											for model_class in self.all_model_classes:
																__lowerCAmelCase           =	model_class(config=_A      )
																for name, param in model.named_parameters():
																					if "embeddings" not in name and param.requires_grad:
																										self.assertIn(
																										    ((param.data.mean() * 1E9).round() / 1E9).item()		,      [0.0, 1.0]		,      msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""		,      )






@require_vision
@require_torch
class       a__   ( unittest.TestCase				):






						@cached_property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return AutoImageProcessor.from_pretrained("microsoft/focalnet-tiny"      ) if is_vision_available() else None






						@slow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	FocalNetForImageClassification.from_pretrained("microsoft/focalnet-tiny"      ).to(_A      )
											__lowerCAmelCase           =	self.default_image_processor

											__lowerCAmelCase           =	Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"      )
											__lowerCAmelCase           =	image_processor(images=_A		,      return_tensors="pt"      ).to(_A      )

											# forward pass
											with torch.no_grad():
																__lowerCAmelCase           =	model(**_A      )

											# verify the logits
											__lowerCAmelCase           =	torch.Size((1, 1_0_0_0)      )
											self.assertEqual(outputs.logits.shape		,      _A      )
											__lowerCAmelCase           =	torch.tensor([0.21_66, -0.43_68, 0.21_91]      ).to(_A      )
											self.assertTrue(torch.allclose(outputs.logits[0, :3]		,      _A		,      atol=1E-4      )      )
											self.assertTrue(outputs.logits.argmax(dim=-1      ).item()		,      2_8_1      )






@require_torch
class       a__   ( snake_case__			,    unittest.TestCase				):
						_a							:   int =   (FocalNetBackbone,) if is_torch_available() else ()
						_a							:   Tuple =   FocalNetConfig

						_a							:   int =   False






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	FocalNetModelTester(self      )










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       a__   :






						def __init__(       self		,      _A		,      _A=1_4		,      _A=7		,      _A=True		,      _A=True		,      _A=False		,      _A=True		,      _A=9_9		,      _A=3_2		,      _A=4		,      _A=4		,      _A=4		,      _A=3_7		,      _A="gelu"		,      _A=0.1		,      _A=0.1		,      _A=5_1_2		,      _A=0.02		,      ):




											"""simple docstring"""



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






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



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

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

											__lowerCAmelCase           =	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=_A		,      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     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



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






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	2_0
											__lowerCAmelCase           =	model_class_name(_A      )

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

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

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

											__lowerCAmelCase           =	model(_A      )

											__lowerCAmelCase           =	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     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	2_0
											__lowerCAmelCase           =	model_class_name(_A      )

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

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

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

											__lowerCAmelCase           =	model(_A		,      attention_mask=_A      )

											__lowerCAmelCase           =	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__   ( snake_case__			,    snake_case__			,    unittest.TestCase				):
						_a							:   str =   (FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
						_a							:   Any =   (FlaxGPTJForCausalLM,) if is_flax_available() else ()






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	FlaxGPTJModelTester(self      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											for model_class_name in self.all_model_classes:
																__lowerCAmelCase       ,							__lowerCAmelCase       ,							__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs()
																self.model_tester.check_use_cache_forward(_A		,      _A		,      _A		,      _A      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											for model_class_name in self.all_model_classes:
																__lowerCAmelCase       ,							__lowerCAmelCase       ,							__lowerCAmelCase           =	self.model_tester.prepare_config_and_inputs()
																self.model_tester.check_use_cache_forward_with_attn_mask(
																    _A		,      _A		,      _A		,      _A      )






						@tooslow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	GPTaTokenizer.from_pretrained("gpt2"		,      pad_token="<|endoftext|>"		,      padding_side="left"      )
											__lowerCAmelCase           =	tokenizer(["Hello this is a long string", "Hey"]		,      return_tensors="np"		,      padding=_A		,      truncation=_A      )

											__lowerCAmelCase           =	FlaxGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"      )
											__lowerCAmelCase           =	False
											__lowerCAmelCase           =	model.config.eos_token_id

											__lowerCAmelCase           =	jax.jit(model.generate      )

											__lowerCAmelCase           =	jit_generate(
											    inputs["input_ids"]		,      attention_mask=inputs["attention_mask"]		,      pad_token_id=tokenizer.pad_token_id      ).sequences

											__lowerCAmelCase           =	tokenizer.batch_decode(_A		,      skip_special_tokens=_A      )

											__lowerCAmelCase           =	[
											    "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(_A		,      _A      )






						@is_pt_flax_cross_test
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



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

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

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

																					__lowerCAmelCase           =	convert_pytorch_state_dict_to_flax(pt_model.state_dict()		,      _A      )
																					__lowerCAmelCase           =	fx_state

																					with torch.no_grad():
																										__lowerCAmelCase           =	pt_model(**_A      ).to_tuple()

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

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

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






						@is_pt_flax_cross_test
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



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

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

																					__lowerCAmelCase           =	pt_model_class(_A      ).eval()
																					__lowerCAmelCase           =	model_class(_A		,      dtype=jnp.floataa      )

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

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

																					with torch.no_grad():
																										__lowerCAmelCase           =	pt_model(**_A      ).to_tuple()

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

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

																					with torch.no_grad():
																										__lowerCAmelCase           =	pt_model_loaded(**_A      ).to_tuple()

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






						@tooslow
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											for model_class_name in self.all_model_classes:
																__lowerCAmelCase           =	model_class_name.from_pretrained("EleutherAI/gpt-j-6B"      )
																__lowerCAmelCase           =	model(np.ones((1, 1)      )      )
																self.assertIsNotNone(_A      )




