


import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTNeoConfig, 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, 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.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
						import torch







class 			snake_case__:







						'''simple docstring'''




						def __init__(			self				,  __lowercase				,  __lowercase=1_4				,  __lowercase=7				,  __lowercase=True				,  __lowercase=True				,  __lowercase=False				,  __lowercase=True				,  __lowercase=9_9				,  __lowercase=3_2				,  __lowercase=4				,  __lowercase=4				,  __lowercase=[[["global", "local"], 2]]				,  __lowercase=3_7				,  __lowercase="gelu"				,  __lowercase=0.1				,  __lowercase=0.1				,  __lowercase=5_1_2				,  __lowercase=7				,  __lowercase=0.02				,  )   ->      List[Any]:
										lowerCAmelCase_	: str      			=			parent
										lowerCAmelCase_	: int      			=			batch_size
										lowerCAmelCase_	: List[Any]      			=			seq_length
										lowerCAmelCase_	: Union[str, Any]      			=			is_training
										lowerCAmelCase_	: Optional[Any]      			=			use_input_mask
										lowerCAmelCase_	: List[Any]      			=			use_token_type_ids
										lowerCAmelCase_	: List[str]      			=			use_labels
										lowerCAmelCase_	: Union[str, Any]      			=			vocab_size
										lowerCAmelCase_	: Dict      			=			hidden_size
										lowerCAmelCase_	: str      			=			num_hidden_layers
										lowerCAmelCase_	: Dict      			=			num_attention_heads
										lowerCAmelCase_	: Any      			=			attention_types
										lowerCAmelCase_	: List[str]      			=			intermediate_size
										lowerCAmelCase_	: List[str]      			=			hidden_act
										lowerCAmelCase_	: List[Any]      			=			hidden_dropout_prob
										lowerCAmelCase_	: List[str]      			=			attention_probs_dropout_prob
										lowerCAmelCase_	: str      			=			max_position_embeddings
										lowerCAmelCase_	: List[Any]      			=			window_size
										lowerCAmelCase_	: Union[str, Any]      			=			initializer_range
										lowerCAmelCase_	: List[str]      			=			None
										lowerCAmelCase_	: Any      			=			vocab_size - 1
										lowerCAmelCase_	: List[Any]      			=			vocab_size - 1
										lowerCAmelCase_	: Optional[Any]      			=			vocab_size - 1




						def   lowercase_	(			self							)   ->      Any:
										lowerCAmelCase_	: Union[str, Any]      			=			ids_tensor([self.batch_size, self.seq_length]				,  self.vocab_size							)

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

										lowerCAmelCase_	: str      			=			GPTNeoConfig(
										    vocab_size=self.vocab_size				,  hidden_size=self.hidden_size				,  num_layers=self.num_hidden_layers				,  num_heads=self.num_attention_heads				,  max_position_embeddings=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				,  window_size=self.window_size				,  attention_types=self.attention_types				,  )

										return (config, input_ids, input_mask)




						def   lowercase_	(			self							)   ->      int:
										lowerCAmelCase_	: str      			=			self.prepare_config_and_inputs()
										lowerCAmelCase_			,					lowerCAmelCase_			,					lowerCAmelCase_	: Optional[int]      			=			config_and_inputs
										lowerCAmelCase_	: Union[str, Any]      			=			{'''input_ids''': input_ids, '''attention_mask''': attention_mask}
										return config, inputs_dict




						def   lowercase_	(			self				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase							)   ->      Dict:
										lowerCAmelCase_	: Optional[int]      			=			2_0
										lowerCAmelCase_	: str      			=			model_class_name(__lowercase							)

										lowerCAmelCase_	: Optional[Any]      			=			model.init_cache(input_ids.shape[0]				,  __lowercase							)
										lowerCAmelCase_	: Dict      			=			jnp.ones((input_ids.shape[0], max_decoder_length)				,  dtype='''i4'''							)

										lowerCAmelCase_	: Dict      			=			jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1							)[None, :]				,  (input_ids.shape[0], input_ids.shape[-1] - 1)							)
										lowerCAmelCase_	: int      			=			model(
										    input_ids[:, :-1]				,  attention_mask=__lowercase				,  past_key_values=__lowercase				,  position_ids=__lowercase				,  )

										lowerCAmelCase_	: int      			=			jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,  dtype='''i4'''							)
										lowerCAmelCase_	: Union[str, Any]      			=			model(
										    input_ids[:, -1:]				,  attention_mask=__lowercase				,  past_key_values=outputs_cache.past_key_values				,  position_ids=__lowercase				,  )

										lowerCAmelCase_	: Tuple      			=			model(__lowercase							)

										lowerCAmelCase_	: 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}"""							)




						def   lowercase_	(			self				,  __lowercase				,  __lowercase				,  __lowercase				,  __lowercase							)   ->      Optional[int]:
										lowerCAmelCase_	: str      			=			2_0
										lowerCAmelCase_	: Optional[int]      			=			model_class_name(__lowercase							)

										lowerCAmelCase_	: Optional[int]      			=			jnp.concatenate(
										    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])							)]				,  axis=-1				,  )

										lowerCAmelCase_	: Union[str, Any]      			=			model.init_cache(input_ids.shape[0]				,  __lowercase							)
										lowerCAmelCase_	: List[Any]      			=			jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1							)[None, :]				,  (input_ids.shape[0], input_ids.shape[-1] - 1)							)

										lowerCAmelCase_	: Optional[int]      			=			model(
										    input_ids[:, :-1]				,  attention_mask=__lowercase				,  past_key_values=__lowercase				,  position_ids=__lowercase				,  )
										lowerCAmelCase_	: Optional[Any]      			=			jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,  dtype='''i4'''							)
										lowerCAmelCase_	: Tuple      			=			model(
										    input_ids[:, -1:]				,  past_key_values=outputs_cache.past_key_values				,  attention_mask=__lowercase				,  position_ids=__lowercase				,  )

										lowerCAmelCase_	: Optional[Any]      			=			model(__lowercase				,  attention_mask=__lowercase							)

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







@require_flax
class 			snake_case__(							UpperCAmelCase__,  UpperCAmelCase__,  unittest.TestCase							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__ :					Optional[Any]   =							(FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
						SCREAMING_SNAKE_CASE__ :					List[Any]   =							(FlaxGPTNeoForCausalLM,) if is_flax_available() else ()




						def   lowercase_	(			self							)   ->      int:
										lowerCAmelCase_	: Any      			=			FlaxGPTNeoModelTester(self							)




						def   lowercase_	(			self							)   ->      str:
										for model_class_name in self.all_model_classes:
														lowerCAmelCase_			,					lowerCAmelCase_			,					lowerCAmelCase_	: List[str]      			=			self.model_tester.prepare_config_and_inputs()
														self.model_tester.check_use_cache_forward(__lowercase				,  __lowercase				,  __lowercase				,  __lowercase							)




						def   lowercase_	(			self							)   ->      Union[str, Any]:
										for model_class_name in self.all_model_classes:
														lowerCAmelCase_			,					lowerCAmelCase_			,					lowerCAmelCase_	: Any      			=			self.model_tester.prepare_config_and_inputs()
														self.model_tester.check_use_cache_forward_with_attn_mask(
														    __lowercase				,  __lowercase				,  __lowercase				,  __lowercase							)




						@slow
						def   lowercase_	(			self							)   ->      Tuple:
										lowerCAmelCase_	: Tuple      			=			GPTaTokenizer.from_pretrained('''gpt2'''				,  pad_token='''<|endoftext|>'''				,  padding_side='''left'''							)
										lowerCAmelCase_	: str      			=			tokenizer(['''Hello this is a long string''', '''Hey''']				,  return_tensors='''np'''				,  padding=__lowercase				,  truncation=__lowercase							)

										lowerCAmelCase_	: Dict      			=			FlaxGPTNeoForCausalLM.from_pretrained('''EleutherAI/gpt-neo-125M'''							)
										lowerCAmelCase_	: Dict      			=			False
										lowerCAmelCase_	: Optional[Any]      			=			model.config.eos_token_id

										lowerCAmelCase_	: Optional[int]      			=			jax.jit(model.generate							)

										lowerCAmelCase_	: Union[str, Any]      			=			jit_generate(
										    inputs['''input_ids''']				,  attention_mask=inputs['''attention_mask''']				,  pad_token_id=tokenizer.pad_token_id							).sequences

										lowerCAmelCase_	: str      			=			tokenizer.batch_decode(__lowercase				,  skip_special_tokens=__lowercase							)

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

										self.assertListEqual(__lowercase				,  __lowercase							)




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

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

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

																		lowerCAmelCase_	: Union[str, Any]      			=			convert_pytorch_state_dict_to_flax(pt_model.state_dict()				,  __lowercase							)
																		lowerCAmelCase_	: List[Any]      			=			fx_state

																		with torch.no_grad():
																						lowerCAmelCase_	: List[str]      			=			pt_model(**__lowercase							).to_tuple()

																		lowerCAmelCase_	: List[Any]      			=			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							)
																						lowerCAmelCase_	: Optional[Any]      			=			model_class.from_pretrained(__lowercase				,  from_pt=__lowercase							)

																		lowerCAmelCase_	: Dict      			=			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   lowercase_	(			self							)   ->      List[str]:
										lowerCAmelCase_			,					lowerCAmelCase_	: Dict      			=			self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes:
														with self.subTest(model_class.__name__							):
																		# prepare inputs
																		lowerCAmelCase_	: List[str]      			=			self._prepare_for_class(__lowercase				,  __lowercase							)
																		lowerCAmelCase_	: Optional[int]      			=			{k: torch.tensor(v.tolist()							) for k, v in prepared_inputs_dict.items()}

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

																		lowerCAmelCase_	: List[str]      			=			pt_model_class(__lowercase							).eval()
																		lowerCAmelCase_	: str      			=			model_class(__lowercase				,  dtype=jnp.floataa							)

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

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

																		with torch.no_grad():
																						lowerCAmelCase_	: Optional[int]      			=			pt_model(**__lowercase							).to_tuple()

																		lowerCAmelCase_	: int      			=			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							)
																						lowerCAmelCase_	: Optional[Any]      			=			pt_model_class.from_pretrained(__lowercase				,  from_flax=__lowercase							)

																		with torch.no_grad():
																						lowerCAmelCase_	: List[Any]      			=			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   lowercase_	(			self							)   ->      List[Any]:
										for model_class_name in self.all_model_classes:
														lowerCAmelCase_	: Tuple      			=			model_class_name.from_pretrained('''EleutherAI/gpt-neo-125M'''							)
														lowerCAmelCase_	: List[Any]      			=			model(np.ones((1, 1)							)							)
														self.assertIsNotNone(__lowercase							)



from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
						import torch

_UpperCAmelCase     :					List[Any]       =logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
						from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):



						def      lowerCAmelCase	(						cls	)->	int:
										return cls



def      lowerCAmelCase	(						lowerCAmelCase_	,				lowerCAmelCase_ = None	,				lowerCAmelCase_ = None	,				lowerCAmelCase_ = None	,				lowerCAmelCase_ = None	,				)->	int:
				lowerCAmelCase_	: Tuple      			=			device
				lowerCAmelCase_	: List[str]      			=			shape[0]

				lowerCAmelCase_	: Union[str, Any]      			=			layout or torch.strided
				lowerCAmelCase_	: Union[str, Any]      			=			device or torch.device('''cpu'''	)

				if generator is not None:
								lowerCAmelCase_	: List[str]      			=			generator.device.type if not isinstance(lowerCAmelCase_	,				lowerCAmelCase_	) else generator[0].device.type
								if gen_device_type != device.type and gen_device_type == "cpu":
												lowerCAmelCase_	: Optional[Any]      			=			'''cpu'''
												if device != "mps":
																logger.info(
																    f"""The passed generator was created on 'cpu' even though a tensor on {device} was expected."""
																    f""" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably"""
																    f""" slighly speed up this function by passing a generator that was created on the {device} device."""	)
								elif gen_device_type != device.type and gen_device_type == "cuda":
												raise ValueError(f"""Cannot generate a {device} tensor from a generator of type {gen_device_type}."""	)

				if isinstance(lowerCAmelCase_	,				lowerCAmelCase_	):
								lowerCAmelCase_	: Dict      			=			(1,) + shape[1:]
								lowerCAmelCase_	: Union[str, Any]      			=			[
								    torch.randn(lowerCAmelCase_	,				generator=generator[i]	,				device=lowerCAmelCase_	,				dtype=lowerCAmelCase_	,				layout=lowerCAmelCase_	)
								    for i in range(lowerCAmelCase_	)
								]
								lowerCAmelCase_	: Tuple      			=			torch.cat(lowerCAmelCase_	,				dim=0	).to(lowerCAmelCase_	)
				else:
								lowerCAmelCase_	: int      			=			torch.randn(lowerCAmelCase_	,				generator=lowerCAmelCase_	,				device=lowerCAmelCase_	,				dtype=lowerCAmelCase_	,				layout=lowerCAmelCase_	).to(lowerCAmelCase_	)

				return latents




def      lowerCAmelCase	(						lowerCAmelCase_	)->	Optional[Any]:
				if is_torch_version('''<'''	,				'''2.0.0'''	) or not hasattr(lowerCAmelCase_	,				'''_dynamo'''	):
								return False
				return isinstance(lowerCAmelCase_	,				torch._dynamo.eval_frame.OptimizedModule	)