





import inspect
import os
from argparse import ArgumentParser, Namespace
from importlib import import_module

import huggingface_hub
import numpy as np
from packaging import version

from .. import (
    FEATURE_EXTRACTOR_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    PROCESSOR_MAPPING,
    TOKENIZER_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    AutoImageProcessor,
    AutoProcessor,
    AutoTokenizer,
    is_datasets_available,
    is_tf_available,
    is_torch_available,
)
from ..utils import TF2_WEIGHTS_INDEX_NAME, TF2_WEIGHTS_NAME, logging
from . import BaseTransformersCLICommand


if is_tf_available():
				import tensorflow as tf

				tf.config.experimental.enable_tensor_float_aa_execution(False)

if is_torch_available():
				import torch

if is_datasets_available():
				from datasets import load_dataset


snake_case_				=      5E-5  # larger error tolerance than in our internal tests, to avoid flaky user-facing errors


def        A__						(							SCREAMING_SNAKE_CASE_		) ->  Tuple:
		return PTtoTFCommand(
		    args.model_name				,							args.local_dir				,							args.max_error				,							args.new_weights				,							args.no_pr				,							args.push				,							args.extra_commit_description				,							args.override_model_class				,							)



class 							snake_case_		(      _A):







		@staticmethod
		def       __lowercase (		__lowercase )					-> int:
				lowerCamelCase :			Any    		=parser.add_parser(
				    '''pt-to-tf'''      ,  help=(
				        '''CLI tool to run convert a transformers model from a PyTorch checkpoint to a TensorFlow checkpoint.'''
				        ''' Can also be used to validate existing weights without opening PRs, with --no-pr.'''
				    )      ,  )
				train_parser.add_argument(
				    '''--model-name'''      ,  type=__lowercase      ,  required=__lowercase      ,  help='''The model name, including owner/organization, as seen on the hub.'''      ,  )
				train_parser.add_argument(
				    '''--local-dir'''      ,  type=__lowercase      ,  default=''''''      ,  help='''Optional local directory of the model repository. Defaults to /tmp/{model_name}'''      ,  )
				train_parser.add_argument(
				    '''--max-error'''      ,  type=__lowercase      ,  default=__lowercase      ,  help=(
				        F"Maximum error tolerance. Defaults to {MAX_ERROR}. This flag should be avoided, use at your own risk."
				    )      ,  )
				train_parser.add_argument(
				    '''--new-weights'''      ,  action='''store_true'''      ,  help='''Optional flag to create new TensorFlow weights, even if they already exist.'''      ,  )
				train_parser.add_argument(
				    '''--no-pr'''      ,  action='''store_true'''      ,  help='''Optional flag to NOT open a PR with converted weights.''' )
				train_parser.add_argument(
				    '''--push'''      ,  action='''store_true'''      ,  help='''Optional flag to push the weights directly to `main` (requires permissions)'''      ,  )
				train_parser.add_argument(
				    '''--extra-commit-description'''      ,  type=__lowercase      ,  default=''''''      ,  help='''Optional additional commit description to use when opening a PR (e.g. to tag the owner).'''      ,  )
				train_parser.add_argument(
				    '''--override-model-class'''      ,  type=__lowercase      ,  default=__lowercase      ,  help='''If you think you know better than the auto-detector, you can specify the model class here. '''
				    '''Can be either an AutoModel class or a specific model class like BertForSequenceClassification.'''      ,  )
				train_parser.set_defaults(func=__lowercase )







		@staticmethod
		def       __lowercase (		__lowercase      ,  __lowercase )					-> Tuple:
				lowerCamelCase :			Tuple    		=set(pt_outputs.keys() )
				lowerCamelCase :			Optional[int]    		=set(tf_outputs.keys() )
				if pt_out_attrs != tf_out_attrs:
						raise ValueError(
						    F"The model outputs have different attributes, aborting. (Pytorch: {pt_out_attrs}, TensorFlow:"
						    F" {tf_out_attrs})" )

				# 2. For each output attribute, computes the difference
				def _find_pt_tf_differences(__lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase="" ):
						# If the current attribute is a tensor, it is a leaf and we make the comparison. Otherwise, we will dig in
						# recursivelly, keeping the name of the attribute.
						if isinstance(__lowercase      ,  torch.Tensor ):
								lowerCamelCase :			Any    		=np.max(np.abs(pt_out.numpy() - tf_out.numpy() ) )
								lowerCamelCase :			int    		=tensor_difference
						else:
								lowerCamelCase :			Dict    		=attr_name
								for i, pt_item in enumerate(__lowercase ):
										# If it is a named attribute, we keep the name. Otherwise, just its index.
										if isinstance(__lowercase      ,  __lowercase ):
												lowerCamelCase :			List[str]    		=root_name + pt_item
												lowerCamelCase :			Tuple    		=tf_out[pt_item]
												lowerCamelCase :			Union[str, Any]    		=pt_out[pt_item]
										else:
												lowerCamelCase :			Tuple    		=root_name + F"[{i}]"
												lowerCamelCase :			Dict    		=tf_out[i]
										lowerCamelCase :			List[str]    		=_find_pt_tf_differences(__lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase )

						return differences

				return _find_pt_tf_differences(__lowercase      ,  __lowercase      ,  {} )







		def __init__(		self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  *__lowercase      ,  )					-> Union[str, Any]:
				lowerCamelCase :			Optional[Any]    		=logging.get_logger('''transformers-cli/pt_to_tf''' )
				lowerCamelCase :			Optional[Any]    		=model_name
				lowerCamelCase :			Optional[Any]    		=local_dir if local_dir else os.path.join('''/tmp'''      ,  __lowercase )
				lowerCamelCase :			Tuple    		=max_error
				lowerCamelCase :			int    		=new_weights
				lowerCamelCase :			Optional[Any]    		=no_pr
				lowerCamelCase :			Union[str, Any]    		=push
				lowerCamelCase :			Optional[Any]    		=extra_commit_description
				lowerCamelCase :			List[Any]    		=override_model_class







		def       __lowercase (		self      ,  __lowercase      ,  __lowercase      ,  __lowercase )					-> List[Any]:

				def _get_audio_input():
						lowerCamelCase :			List[Any]    		=load_dataset('''hf-internal-testing/librispeech_asr_dummy'''      ,  '''clean'''      ,  split='''validation''' )
						lowerCamelCase :			Tuple    		=ds.sort('''id''' ).select(range(2 ) )[:2]['''audio''']
						lowerCamelCase :			str    		=[x['''array'''] for x in speech_samples]
						return raw_samples

				lowerCamelCase :			List[Any]    		=type(pt_model.config )
				if model_config_class in PROCESSOR_MAPPING:
						lowerCamelCase :			Optional[int]    		=AutoProcessor.from_pretrained(self._local_dir )
						if model_config_class in TOKENIZER_MAPPING and processor.tokenizer.pad_token is None:
								lowerCamelCase :			List[str]    		=processor.tokenizer.eos_token
				elif model_config_class in IMAGE_PROCESSOR_MAPPING:
						lowerCamelCase :			List[Any]    		=AutoImageProcessor.from_pretrained(self._local_dir )
				elif model_config_class in FEATURE_EXTRACTOR_MAPPING:
						lowerCamelCase :			Optional[Any]    		=AutoFeatureExtractor.from_pretrained(self._local_dir )
				elif model_config_class in TOKENIZER_MAPPING:
						lowerCamelCase :			Optional[int]    		=AutoTokenizer.from_pretrained(self._local_dir )
						if processor.pad_token is None:
								lowerCamelCase :			Optional[int]    		=processor.eos_token
				else:
						raise ValueError(F"Unknown data processing type (model config type: {model_config_class})" )

				lowerCamelCase :			Optional[int]    		=set(inspect.signature(pt_model.forward ).parameters.keys() )
				lowerCamelCase :			Dict    		={}
				if "input_ids" in model_forward_signature:
						processor_inputs.update(
						    {
						        '''text''': ['''Hi there!''', '''I am a batch with more than one row and different input lengths.'''],
						        '''padding''': True,
						        '''truncation''': True,
						    } )
				if "pixel_values" in model_forward_signature:
						lowerCamelCase :			Any    		=load_dataset('''cifar10'''      ,  '''plain_text'''      ,  split='''test''' )[:2]['''img''']
						processor_inputs.update({'''images''': sample_images} )
				if "input_features" in model_forward_signature:
						lowerCamelCase :			Dict    		=inspect.signature(processor.feature_extractor ).parameters
						# Pad to the largest input length by default but take feature extractor default
						# padding value if it exists e.g. "max_length" and is not False or None
						if "padding" in feature_extractor_signature:
								lowerCamelCase :			Union[str, Any]    		=feature_extractor_signature['''padding'''].default
								if default_strategy is not False and default_strategy is not None:
										lowerCamelCase :			Optional[int]    		=default_strategy
								else:
										lowerCamelCase :			List[Any]    		=True
						else:
								lowerCamelCase :			Any    		=True
						processor_inputs.update({'''audio''': _get_audio_input(), '''padding''': padding_strategy} )
				if "input_values" in model_forward_signature:  # Wav2Vec2 audio input
						processor_inputs.update({'''audio''': _get_audio_input(), '''padding''': True} )
				lowerCamelCase :			Optional[Any]    		=processor(**__lowercase      ,  return_tensors='''pt''' )
				lowerCamelCase :			str    		=processor(**__lowercase      ,  return_tensors='''tf''' )

				# Extra input requirements, in addition to the input modality
				if (
				    config.is_encoder_decoder
				    or (hasattr(__lowercase      ,  '''encoder''' ) and hasattr(__lowercase      ,  '''decoder''' ))
				    or "decoder_input_ids" in tf_dummy_inputs
				):
						lowerCamelCase :			Optional[int]    		=np.asarray([[1], [1]]      ,  dtype=__lowercase ) * (pt_model.config.decoder_start_token_id or 0)
						pt_input.update({'''decoder_input_ids''': torch.tensor(__lowercase )} )
						tf_input.update({'''decoder_input_ids''': tf.convert_to_tensor(__lowercase )} )

				return pt_input, tf_input







		def       __lowercase (		self )					-> Optional[int]:
				# hub version 0.9.0 introduced the possibility of programmatically opening PRs with normal write tokens.
				if version.parse(huggingface_hub.__version__ ) < version.parse('''0.9.0''' ):
						raise ImportError(
						    '''The huggingface_hub version must be >= 0.9.0 to use this command. Please update your huggingface_hub'''
						    ''' installation.''' )
				else:
						from huggingface_hub import Repository, create_commit
						from huggingface_hub._commit_api import CommitOperationAdd

				# Fetch remote data
				lowerCamelCase :			Union[str, Any]    		=Repository(local_dir=self._local_dir      ,  clone_from=self._model_name )

				# Load config and get the appropriate architecture -- the latter is needed to convert the head's weights
				lowerCamelCase :			int    		=AutoConfig.from_pretrained(self._local_dir )
				lowerCamelCase :			Optional[int]    		=config.architectures
				if self._override_model_class is not None:
						if self._override_model_class.startswith('''TF''' ):
								lowerCamelCase :			str    		=[self._override_model_class[2:]]
						else:
								lowerCamelCase :			int    		=[self._override_model_class]
						try:
								lowerCamelCase :			Any    		=getattr(import_module('''transformers''' )      ,  architectures[0] )
						except AttributeError:
								raise ValueError(F"Model class {self._override_model_class} not found in transformers." )
						try:
								lowerCamelCase :			Union[str, Any]    		=getattr(import_module('''transformers''' )      ,  '''TF''' + architectures[0] )
						except AttributeError:
								raise ValueError(F"TF model class TF{self._override_model_class} not found in transformers." )
				elif architectures is None:  # No architecture defined -- use auto classes
						lowerCamelCase :			Tuple    		=getattr(import_module('''transformers''' )      ,  '''AutoModel''' )
						lowerCamelCase :			Union[str, Any]    		=getattr(import_module('''transformers''' )      ,  '''TFAutoModel''' )
						self._logger.warning('''No detected architecture, using AutoModel/TFAutoModel''' )
				else:  # Architecture defined -- use it
						if len(__lowercase ) > 1:
								raise ValueError(F"More than one architecture was found, aborting. (architectures = {architectures})" )
						self._logger.warning(F"Detected architecture: {architectures[0]}" )
						lowerCamelCase :			List[Any]    		=getattr(import_module('''transformers''' )      ,  architectures[0] )
						try:
								lowerCamelCase :			Optional[int]    		=getattr(import_module('''transformers''' )      ,  '''TF''' + architectures[0] )
						except AttributeError:
								raise AttributeError(F"The TensorFlow equivalent of {architectures[0]} doesn't exist in transformers." )

        # Check the TF dummy inputs to see what keys we need in the forward pass
				lowerCamelCase :			Dict    		=tf_class.from_config(__lowercase )
				lowerCamelCase :			Union[str, Any]    		=tf_from_pt_model.dummy_inputs

				del tf_from_pt_model  # Try to keep only one model in memory at a time

				# Load the model and get some basic inputs
				lowerCamelCase :			Tuple    		=pt_class.from_pretrained(self._local_dir )
				pt_model.eval()

				lowerCamelCase		, lowerCamelCase :			Tuple    		=self.get_inputs(__lowercase      ,  __lowercase      ,  __lowercase )

				with torch.no_grad():
						lowerCamelCase :			Optional[Any]    		=pt_model(**__lowercase      ,  output_hidden_states=__lowercase )
				del pt_model  # will no longer be used, and may have a large memory footprint

				lowerCamelCase :			List[Any]    		=tf_class.from_pretrained(self._local_dir      ,  from_pt=__lowercase )
				lowerCamelCase :			Optional[Any]    		=tf_from_pt_model(**__lowercase      ,  output_hidden_states=__lowercase      ,  training=__lowercase )

				# Confirms that cross loading PT weights into TF worked.
				lowerCamelCase :			str    		=self.find_pt_tf_differences(__lowercase      ,  __lowercase )
				lowerCamelCase :			int    		={k: v for k, v in crossload_differences.items() if '''hidden''' not in k}
				lowerCamelCase :			Optional[Any]    		={k: v for k, v in crossload_differences.items() if '''hidden''' in k}
				if len(__lowercase ) == 0 and architectures is not None:
						raise ValueError(
						    F"Something went wrong -- the config file has architectures ({architectures}), but no model head"
						    ''' output was found. All outputs start with \'hidden\'''' )
				lowerCamelCase :			str    		=max(output_differences.values() ) if output_differences else 0.0
				lowerCamelCase :			Optional[int]    		=max(hidden_differences.values() )
				if max_crossload_output_diff > self._max_error or max_crossload_hidden_diff > self._max_error:
						raise ValueError(
						    '''The cross-loaded TensorFlow model has different outputs, something went wrong!\n'''
						    + F"\nList of maximum output differences above the threshold ({self._max_error}):\n"
						    + '''\n'''.join([F"{k}: {v:.3e}" for k, v in output_differences.items() if v > self._max_error] )
						    + F"\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n"
						    + '''\n'''.join([F"{k}: {v:.3e}" for k, v in hidden_differences.items() if v > self._max_error] ) )

				# Save the weights in a TF format (if needed) and confirms that the results are still good
				lowerCamelCase :			Optional[int]    		=os.path.join(self._local_dir      ,  __lowercase )
				lowerCamelCase :			Tuple    		=os.path.join(self._local_dir      ,  __lowercase )
				if (not os.path.exists(__lowercase ) and not os.path.exists(__lowercase )) or self._new_weights:
						tf_from_pt_model.save_pretrained(self._local_dir )
				del tf_from_pt_model  # will no longer be used, and may have a large memory footprint

				lowerCamelCase :			int    		=tf_class.from_pretrained(self._local_dir )
				lowerCamelCase :			Dict    		=tf_model(**__lowercase      ,  output_hidden_states=__lowercase )

				lowerCamelCase :			Optional[Any]    		=self.find_pt_tf_differences(__lowercase      ,  __lowercase )
				lowerCamelCase :			Optional[int]    		={k: v for k, v in conversion_differences.items() if '''hidden''' not in k}
				lowerCamelCase :			Any    		={k: v for k, v in conversion_differences.items() if '''hidden''' in k}
				if len(__lowercase ) == 0 and architectures is not None:
						raise ValueError(
						    F"Something went wrong -- the config file has architectures ({architectures}), but no model head"
						    ''' output was found. All outputs start with \'hidden\'''' )
				lowerCamelCase :			Optional[Any]    		=max(output_differences.values() ) if output_differences else 0.0
				lowerCamelCase :			Any    		=max(hidden_differences.values() )
				if max_conversion_output_diff > self._max_error or max_conversion_hidden_diff > self._max_error:
						raise ValueError(
						    '''The converted TensorFlow model has different outputs, something went wrong!\n'''
						    + F"\nList of maximum output differences above the threshold ({self._max_error}):\n"
						    + '''\n'''.join([F"{k}: {v:.3e}" for k, v in output_differences.items() if v > self._max_error] )
						    + F"\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n"
						    + '''\n'''.join([F"{k}: {v:.3e}" for k, v in hidden_differences.items() if v > self._max_error] ) )

				lowerCamelCase :			int    		='''Update TF weights''' if self._new_weights else '''Add TF weights'''
				if self._push:
						repo.git_add(auto_lfs_track=__lowercase )
						repo.git_commit(__lowercase )
						repo.git_push(blocking=__lowercase )  # this prints a progress bar with the upload
						self._logger.warning(F"TF weights pushed into {self._model_name}" )
				elif not self._no_pr:
						self._logger.warning('''Uploading the weights into a new PR...''' )
						lowerCamelCase :			Optional[Any]    		=(
						    '''Model converted by the [`transformers`\' `pt_to_tf`'''
						    ''' CLI](https://github.com/huggingface/transformers/blob/main/src/transformers/commands/pt_to_tf.py). '''
						    '''All converted model outputs and hidden layers were validated against its PyTorch counterpart.\n\n'''
						    F"Maximum crossload output difference={max_crossload_output_diff:.3e}; "
						    F"Maximum crossload hidden layer difference={max_crossload_hidden_diff:.3e};\n"
						    F"Maximum conversion output difference={max_conversion_output_diff:.3e}; "
						    F"Maximum conversion hidden layer difference={max_conversion_hidden_diff:.3e};\n"
						)
						if self._max_error > MAX_ERROR:
								commit_descrition += (
								    F"\n\nCAUTION: The maximum admissible error was manually increased to {self._max_error}!"
								)
						if self._extra_commit_description:
								commit_descrition += "\n\n" + self._extra_commit_description

						# sharded model -> adds all related files (index and .h5 shards)
						if os.path.exists(__lowercase ):
								lowerCamelCase :			Tuple    		=[
								    CommitOperationAdd(path_in_repo=__lowercase      ,  path_or_fileobj=__lowercase )
								]
								for shard_path in tf.io.gfile.glob(self._local_dir + '''/tf_model-*.h5''' ):
										operations += [
										    CommitOperationAdd(path_in_repo=os.path.basename(__lowercase )      ,  path_or_fileobj=__lowercase )
										]
						else:
								lowerCamelCase :			List[str]    		=[CommitOperationAdd(path_in_repo=__lowercase      ,  path_or_fileobj=__lowercase )]

						lowerCamelCase :			Optional[Any]    		=create_commit(
						    repo_id=self._model_name      ,  operations=__lowercase      ,  commit_message=__lowercase      ,  commit_description=__lowercase      ,  repo_type='''model'''      ,  create_pr=__lowercase      ,  ).pr_url
						self._logger.warning(F"PR open in {hub_pr_url}" )








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







		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.0_2      ,  )					-> Tuple:
				lowerCamelCase :			Optional[Any]    		=parent
				lowerCamelCase :			Optional[int]    		=batch_size
				lowerCamelCase :			Any    		=seq_length
				lowerCamelCase :			Optional[Any]    		=is_training
				lowerCamelCase :			Any    		=use_input_mask
				lowerCamelCase :			int    		=use_token_type_ids
				lowerCamelCase :			int    		=use_labels
				lowerCamelCase :			str    		=vocab_size
				lowerCamelCase :			Optional[int]    		=hidden_size
				lowerCamelCase :			int    		=num_hidden_layers
				lowerCamelCase :			List[Any]    		=num_attention_heads
				lowerCamelCase :			Any    		=attention_types
				lowerCamelCase :			List[str]    		=intermediate_size
				lowerCamelCase :			Dict    		=hidden_act
				lowerCamelCase :			Union[str, Any]    		=hidden_dropout_prob
				lowerCamelCase :			Optional[int]    		=attention_probs_dropout_prob
				lowerCamelCase :			Any    		=max_position_embeddings
				lowerCamelCase :			int    		=window_size
				lowerCamelCase :			List[Any]    		=initializer_range
				lowerCamelCase :			Any    		=None
				lowerCamelCase :			List[Any]    		=vocab_size - 1
				lowerCamelCase :			str    		=vocab_size - 1
				lowerCamelCase :			int    		=vocab_size - 1







		def       __lowercase (		self )					-> Dict:
				lowerCamelCase :			List[Any]    		=ids_tensor([self.batch_size, self.seq_length]      ,  self.vocab_size )

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

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







		def       __lowercase (		self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase )					-> List[Any]:
				lowerCamelCase :			Tuple    		=2_0
				lowerCamelCase :			Union[str, Any]    		=model_class_name(__lowercase )

				lowerCamelCase :			List[Any]    		=model.init_cache(input_ids.shape[0]      ,  __lowercase )
				lowerCamelCase :			str    		=jnp.ones((input_ids.shape[0], max_decoder_length)      ,  dtype='''i4''' )

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

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

				lowerCamelCase :			Any    		=model(__lowercase )

				lowerCamelCase :			Dict    		=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 )					-> Dict:
				lowerCamelCase :			str    		=2_0
				lowerCamelCase :			Optional[int]    		=model_class_name(__lowercase )

				lowerCamelCase :			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[Any]    		=model(
				    input_ids[:, :-1]      ,  attention_mask=__lowercase      ,  past_key_values=__lowercase      ,  position_ids=__lowercase      ,  )
				lowerCamelCase :			List[Any]    		=jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]      ,  dtype='''i4''' )
				lowerCamelCase :			Union[str, Any]    		=model(
				    input_ids[:, -1:]      ,  past_key_values=outputs_cache.past_key_values      ,  attention_mask=__lowercase      ,  position_ids=__lowercase      ,  )

				lowerCamelCase :			Union[str, Any]    		=model(__lowercase      ,  attention_mask=__lowercase )

				lowerCamelCase :			Dict    		=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_		(      _A    ,       _A    ,       unittest.TestCase):
		lowerCamelCase     :List[Any]							=  (FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
		lowerCamelCase     :Union[str, Any]							=  (FlaxGPTNeoForCausalLM,) if is_flax_available() else ()







		def       __lowercase (		self )					-> Union[str, Any]:
				lowerCamelCase :			str    		=FlaxGPTNeoModelTester(self )







		def       __lowercase (		self )					-> Dict:
				for model_class_name in self.all_model_classes:
						lowerCamelCase		, lowerCamelCase		, lowerCamelCase :			Dict    		=self.model_tester.prepare_config_and_inputs()
						self.model_tester.check_use_cache_forward(__lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase )







		def       __lowercase (		self )					-> Tuple:
				for model_class_name in self.all_model_classes:
						lowerCamelCase		, lowerCamelCase		, lowerCamelCase :			Optional[int]    		=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 )					-> Optional[int]:
				lowerCamelCase :			Tuple    		=GPTaTokenizer.from_pretrained('''gpt2'''      ,  pad_token='''<|endoftext|>'''      ,  padding_side='''left''' )
				lowerCamelCase :			Any    		=tokenizer(['''Hello this is a long string''', '''Hey''']      ,  return_tensors='''np'''      ,  padding=__lowercase      ,  truncation=__lowercase )

				lowerCamelCase :			List[Any]    		=FlaxGPTNeoForCausalLM.from_pretrained('''EleutherAI/gpt-neo-125M''' )
				lowerCamelCase :			Union[str, Any]    		=False
				lowerCamelCase :			Optional[int]    		=model.config.eos_token_id

				lowerCamelCase :			str    		=jax.jit(model.generate )

				lowerCamelCase :			Optional[int]    		=jit_generate(
				    inputs['''input_ids''']      ,  attention_mask=inputs['''attention_mask''']      ,  pad_token_id=tokenizer.pad_token_id ).sequences

				lowerCamelCase :			Any    		=tokenizer.batch_decode(__lowercase      ,  skip_special_tokens=__lowercase )

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

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

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

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

								with torch.no_grad():
										lowerCamelCase :			List[Any]    		=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 :			List[str]    		=model_class.from_pretrained(__lowercase      ,  from_pt=__lowercase )

								lowerCamelCase :			List[str]    		=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 )					-> 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 :			Tuple    		=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 :			Optional[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 :			Union[str, Any]    		=model_class(__lowercase      ,  dtype=jnp.floataa )

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

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

								with torch.no_grad():
										lowerCamelCase :			Tuple    		=pt_model(**__lowercase ).to_tuple()

								lowerCamelCase :			Optional[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:
										fx_model.save_pretrained(__lowercase )
										lowerCamelCase :			Union[str, Any]    		=pt_model_class.from_pretrained(__lowercase      ,  from_flax=__lowercase )

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

