






"""simple docstring"""






import sacrebleu as scb
from packaging import version
from sacrebleu import TER

import datasets


UpperCAmelCase_     :				Optional[Any]         =				"""\
@inproceedings{snover-etal-2006-study,
    title = \"A Study of Translation Edit Rate with Targeted Human Annotation\",
    author = \"Snover, Matthew  and
      Dorr, Bonnie  and
      Schwartz, Rich  and
      Micciulla, Linnea  and
      Makhoul, John\",
    booktitle = \"Proceedings of the 7th Conference of the Association for Machine Translation in the Americas: Technical Papers\",
    month = aug # \" 8-12\",
    year = \"2006\",
    address = \"Cambridge, Massachusetts, USA\",
    publisher = \"Association for Machine Translation in the Americas\",
    url = \"https://aclanthology.org/2006.amta-papers.25\",
    pages = \"223--231\",
}
@inproceedings{post-2018-call,
    title = \"A Call for Clarity in Reporting {BLEU} Scores\",
    author = \"Post, Matt\",
    booktitle = \"Proceedings of the Third Conference on Machine Translation: Research Papers\",
    month = oct,
    year = \"2018\",
    address = \"Belgium, Brussels\",
    publisher = \"Association for Computational Linguistics\",
    url = \"https://www.aclweb.org/anthology/W18-6319\",
    pages = \"186--191\",
}
"""

UpperCAmelCase_     :				Any         =				"""\
TER (Translation Edit Rate, also called Translation Error Rate) is a metric to quantify the edit operations that a
hypothesis requires to match a reference translation. We use the implementation that is already present in sacrebleu
(https://github.com/mjpost/sacreBLEU#ter), which in turn is inspired by the TERCOM implementation, which can be found
here: https://github.com/jhclark/tercom.

The implementation here is slightly different from sacrebleu in terms of the required input format. The length of
the references and hypotheses lists need to be the same, so you may need to transpose your references compared to
sacrebleu's required input format. See https://github.com/huggingface/datasets/issues/3154#issuecomment-950746534

See the README.md file at https://github.com/mjpost/sacreBLEU#ter for more information.
"""

UpperCAmelCase_     :				Union[str, Any]         =				"""
Produces TER scores alongside the number of edits and reference length.

Args:
    predictions (list of str): The system stream (a sequence of segments).
    references (list of list of str): A list of one or more reference streams (each a sequence of segments).
    normalized (boolean): If `True`, applies basic tokenization and normalization to sentences. Defaults to `False`.
    ignore_punct (boolean): If `True`, applies basic tokenization and normalization to sentences. Defaults to `False`.
    support_zh_ja_chars (boolean): If `True`, tokenization/normalization supports processing of Chinese characters,
                                    as well as Japanese Kanji, Hiragana, Katakana, and Phonetic Extensions of Katakana.
                                    Only applies if `normalized = True`. Defaults to `False`.
    case_sensitive (boolean): If `False`, makes all predictions and references lowercase to ignore differences in case. Defaults to `False`.

Returns:
    'score' (float): TER score (num_edits / sum_ref_lengths * 100)
    'num_edits' (int): The cumulative number of edits
    'ref_length' (float): The cumulative average reference length

Examples:
    Example 1:
        >>> predictions = [\"does this sentence match??\",
        ...                     \"what about this sentence?\",
        ...                     \"What did the TER metric user say to the developer?\"]
        >>> references = [[\"does this sentence match\", \"does this sentence match!?!\"],
        ...             [\"wHaT aBoUt ThIs SeNtEnCe?\", \"wHaT aBoUt ThIs SeNtEnCe?\"],
        ...             [\"Your jokes are...\", \"...TERrible\"]]
        >>> ter = datasets.load_metric(\"ter\")
        >>> results = ter.compute(predictions=predictions,
        ...                         references=references,
        ...                         case_sensitive=True)
        >>> print(results)
        {'score': 150.0, 'num_edits': 15, 'ref_length': 10.0}

    Example 2:
        >>> predictions = [\"does this sentence match??\",
        ...                     \"what about this sentence?\"]
        >>> references = [[\"does this sentence match\", \"does this sentence match!?!\"],
        ...             [\"wHaT aBoUt ThIs SeNtEnCe?\", \"wHaT aBoUt ThIs SeNtEnCe?\"]]
        >>> ter = datasets.load_metric(\"ter\")
        >>> results = ter.compute(predictions=predictions,
        ...                         references=references,
        ...                         case_sensitive=True)
        >>> print(results)
        {'score': 62.5, 'num_edits': 5, 'ref_length': 8.0}

    Example 3:
        >>> predictions = [\"does this sentence match??\",
        ...                     \"what about this sentence?\"]
        >>> references = [[\"does this sentence match\", \"does this sentence match!?!\"],
        ...             [\"wHaT aBoUt ThIs SeNtEnCe?\", \"wHaT aBoUt ThIs SeNtEnCe?\"]]
        >>> ter = datasets.load_metric(\"ter\")
        >>> results = ter.compute(predictions=predictions,
        ...                         references=references,
        ...                         normalized=True,
        ...                         case_sensitive=True)
        >>> print(results)
        {'score': 57.14285714285714, 'num_edits': 6, 'ref_length': 10.5}

    Example 4:
        >>> predictions = [\"does this sentence match??\",
        ...                     \"what about this sentence?\"]
        >>> references = [[\"does this sentence match\", \"does this sentence match!?!\"],
        ...             [\"wHaT aBoUt ThIs SeNtEnCe?\", \"wHaT aBoUt ThIs SeNtEnCe?\"]]
        >>> ter = datasets.load_metric(\"ter\")
        >>> results = ter.compute(predictions=predictions,
        ...                         references=references,
        ...                         ignore_punct=True,
        ...                         case_sensitive=False)
        >>> print(results)
        {'score': 0.0, 'num_edits': 0, 'ref_length': 8.0}

    Example 5:
        >>> predictions = [\"does this sentence match??\",
        ...                    \"what about this sentence?\",
        ...                    \"What did the TER metric user say to the developer?\"]
        >>> references = [[\"does this sentence match\", \"does this sentence match!?!\"],
        ...             [\"wHaT aBoUt ThIs SeNtEnCe?\", \"wHaT aBoUt ThIs SeNtEnCe?\"],
        ...             [\"Your jokes are...\", \"...TERrible\"]]
        >>> ter = datasets.load_metric(\"ter\")
        >>> results = ter.compute(predictions=predictions,
        ...                         references=references,
        ...                         ignore_punct=True,
        ...                         case_sensitive=False)
        >>> print(results)
        {'score': 100.0, 'num_edits': 10, 'ref_length': 10.0}
"""




@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION       ,	_KWARGS_DESCRIPTION				)
class 							lowerCAmelCase__       (    datasets.Metric				):






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          if version.parse(scb.__version__) < version.parse('''1.4.12'''):
              raise ImportWarning(
                  '''To use `sacrebleu`, the module `sacrebleu>=1.4.12` is required, and the current version of `sacrebleu` doesn\'t match this condition.\n'''
                  '''You can install it with `pip install "sacrebleu>=1.4.12"`.''')
          return datasets.MetricInfo(
              description=_DESCRIPTION		,			citation=_CITATION		,			homepage='''http://www.cs.umd.edu/~snover/tercom/'''		,			inputs_description=_KWARGS_DESCRIPTION		,			features=datasets.Features(
                  {
                      '''predictions''': datasets.Value('''string'''		,			id='''sequence'''),
                      '''references''': datasets.Sequence(datasets.Value('''string'''		,			id='''sequence''')		,			id='''references'''),
                  })		,			codebase_urls=['''https://github.com/mjpost/sacreBLEU#ter''']		,			reference_urls=[
                  '''https://github.com/jhclark/tercom''',
              ]		,			)



      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]		,			lowercase_  :						Optional[Any]		,			lowercase_  :						Tuple		,			lowercase_  :						bool = False		,			lowercase_  :						bool = False		,			lowercase_  :						bool = False		,			lowercase_  :						bool = False		,			):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	len(references[0])
          if any(len(lowercase_) != references_per_prediction for refs in references):
              raise ValueError('''Sacrebleu requires the same number of references for each prediction''')
          SCREAMING_SNAKE_CASE_   :						Tuple							=	[[refs[i] for refs in references] for i in range(lowercase_)]

          SCREAMING_SNAKE_CASE_   :						Tuple							=	TER(
              normalized=lowercase_		,			no_punct=lowercase_		,			asian_support=lowercase_		,			case_sensitive=lowercase_		,			)
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	sb_ter.corpus_score(lowercase_		,			lowercase_)

          return {"score": output.score, "num_edits": output.num_edits, "ref_length": output.ref_length}










"""simple docstring"""






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






      '''simple docstring'''

      def __init__(		self  :						int		,			lowercase_  :						List[Any]		,			lowercase_  :						int=14		,			lowercase_  :						Tuple=7		,			lowercase_  :						int=True		,			lowercase_  :						int=True		,			lowercase_  :						int=False		,			lowercase_  :						Optional[int]=True		,			lowercase_  :						Union[str, Any]=99		,			lowercase_  :						Optional[int]=32		,			lowercase_  :						List[Any]=4		,			lowercase_  :						Optional[int]=4		,			lowercase_  :						List[str]=4		,			lowercase_  :						Tuple=37		,			lowercase_  :						Any="gelu"		,			lowercase_  :						Dict=0.1		,			lowercase_  :						str=0.1		,			lowercase_  :						Optional[int]=512		,			lowercase_  :						Dict=0.02		,			):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	parent
          SCREAMING_SNAKE_CASE_   :						Any							=	batch_size
          SCREAMING_SNAKE_CASE_   :						Tuple							=	seq_length
          SCREAMING_SNAKE_CASE_   :						Any							=	is_training
          SCREAMING_SNAKE_CASE_   :						int							=	use_input_mask
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	use_token_type_ids
          SCREAMING_SNAKE_CASE_   :						Tuple							=	use_labels
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	vocab_size
          SCREAMING_SNAKE_CASE_   :						str							=	hidden_size
          SCREAMING_SNAKE_CASE_   :						Tuple							=	rotary_dim
          SCREAMING_SNAKE_CASE_   :						str							=	num_hidden_layers
          SCREAMING_SNAKE_CASE_   :						int							=	num_attention_heads
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	intermediate_size
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	hidden_act
          SCREAMING_SNAKE_CASE_   :						str							=	hidden_dropout_prob
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	attention_probs_dropout_prob
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	max_position_embeddings
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	initializer_range
          SCREAMING_SNAKE_CASE_   :						Dict							=	None
          SCREAMING_SNAKE_CASE_   :						str							=	vocab_size - 1
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	vocab_size - 1
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	vocab_size - 1
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	ids_tensor([self.batch_size, self.seq_length]		,			self.vocab_size)

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

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	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=lowercase_		,			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  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	self.prepare_config_and_inputs()
          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Tuple							=	config_and_inputs
          SCREAMING_SNAKE_CASE_   :						int							=	{'''input_ids''': input_ids, '''attention_mask''': attention_mask}
          return config, inputs_dict
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple		,			lowercase_  :						Tuple		,			lowercase_  :						Any		,			lowercase_  :						Any		,			lowercase_  :						Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	20
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	model_class_name(lowercase_)

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

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

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

          SCREAMING_SNAKE_CASE_   :						List[str]							=	model(lowercase_)

          SCREAMING_SNAKE_CASE_   :						int							=	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  :						Tuple		,			lowercase_  :						Tuple		,			lowercase_  :						Tuple		,			lowercase_  :						Optional[int]		,			lowercase_  :						Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	20
          SCREAMING_SNAKE_CASE_   :						List[str]							=	model_class_name(lowercase_)

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

          SCREAMING_SNAKE_CASE_   :						Tuple							=	model.init_cache(input_ids.shape[0]		,			lowercase_)
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	jnp.broadcast_to(
              jnp.arange(input_ids.shape[-1] - 1)[None, :]		,			(input_ids.shape[0], input_ids.shape[-1] - 1))

          SCREAMING_SNAKE_CASE_   :						Dict							=	model(
              input_ids[:, :-1]		,			attention_mask=lowercase_		,			past_key_values=lowercase_		,			position_ids=lowercase_		,			)
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]		,			dtype='''i4''')
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	model(
              input_ids[:, -1:]		,			past_key_values=outputs_cache.past_key_values		,			attention_mask=lowercase_		,			position_ids=lowercase_		,			)

          SCREAMING_SNAKE_CASE_   :						Any							=	model(lowercase_		,			attention_mask=lowercase_)

          SCREAMING_SNAKE_CASE_   :						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 							lowerCAmelCase__       (    UpperCAmelCase__       ,	UpperCAmelCase__       ,	unittest.TestCase				):






      '''simple docstring'''

      __UpperCamelCase							=						(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
      __UpperCamelCase							=						(FlaxGPTJForCausalLM,) if is_flax_available() else ()
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[str]							=	FlaxGPTJModelTester(self)
      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          for model_class_name in self.all_model_classes:
              SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						int							=	self.model_tester.prepare_config_and_inputs()
              self.model_tester.check_use_cache_forward(lowercase_		,			lowercase_		,			lowercase_		,			lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          for model_class_name in self.all_model_classes:
              SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Optional[int]							=	self.model_tester.prepare_config_and_inputs()
              self.model_tester.check_use_cache_forward_with_attn_mask(
                  lowercase_		,			lowercase_		,			lowercase_		,			lowercase_)
      @tooslow
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	GPTaTokenizer.from_pretrained('''gpt2'''		,			pad_token='''<|endoftext|>'''		,			padding_side='''left''')
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	tokenizer(['''Hello this is a long string''', '''Hey''']		,			return_tensors='''np'''		,			padding=lowercase_		,			truncation=lowercase_)

          SCREAMING_SNAKE_CASE_   :						List[str]							=	FlaxGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B''')
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	False
          SCREAMING_SNAKE_CASE_   :						str							=	model.config.eos_token_id

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	jax.jit(model.generate)

          SCREAMING_SNAKE_CASE_   :						Tuple							=	jit_generate(
              inputs['''input_ids''']		,			attention_mask=inputs['''attention_mask''']		,			pad_token_id=tokenizer.pad_token_id).sequences

          SCREAMING_SNAKE_CASE_   :						int							=	tokenizer.batch_decode(lowercase_		,			skip_special_tokens=lowercase_)

          SCREAMING_SNAKE_CASE_   :						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 _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Optional[int]							=	self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__):
                  # prepare inputs
                  SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	self._prepare_for_class(lowercase_		,			lowercase_)
                  SCREAMING_SNAKE_CASE_   :						Any							=	{k: torch.tensor(v.tolist()) for k, v in prepared_inputs_dict.items()}

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

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

                  SCREAMING_SNAKE_CASE_   :						List[str]							=	convert_pytorch_state_dict_to_flax(pt_model.state_dict()		,			lowercase_)
                  SCREAMING_SNAKE_CASE_   :						List[str]							=	fx_state

                  with torch.no_grad():
                      SCREAMING_SNAKE_CASE_   :						Dict							=	pt_model(**lowercase_).to_tuple()

                  SCREAMING_SNAKE_CASE_   :						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_)
                      SCREAMING_SNAKE_CASE_   :						int							=	model_class.from_pretrained(lowercase_		,			from_pt=lowercase_)

                  SCREAMING_SNAKE_CASE_   :						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 _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						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
                  SCREAMING_SNAKE_CASE_   :						Tuple							=	self._prepare_for_class(lowercase_		,			lowercase_)
                  SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	{k: torch.tensor(v.tolist()) for k, v in prepared_inputs_dict.items()}

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

                  SCREAMING_SNAKE_CASE_   :						int							=	pt_model_class(lowercase_).eval()
                  SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	model_class(lowercase_		,			dtype=jnp.floataa)

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

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

                  with torch.no_grad():
                      SCREAMING_SNAKE_CASE_   :						Optional[int]							=	pt_model(**lowercase_).to_tuple()

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

                  with tempfile.TemporaryDirectory() as tmpdirname:
                      fx_model.save_pretrained(lowercase_)
                      SCREAMING_SNAKE_CASE_   :						List[str]							=	pt_model_class.from_pretrained(lowercase_		,			from_flax=lowercase_)

                  with torch.no_grad():
                      SCREAMING_SNAKE_CASE_   :						Union[str, 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)



      @tooslow
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]):




          '''simple docstring'''


          for model_class_name in self.all_model_classes:
              SCREAMING_SNAKE_CASE_   :						Any							=	model_class_name.from_pretrained('''EleutherAI/gpt-j-6B''')
              SCREAMING_SNAKE_CASE_   :						List[str]							=	model(np.ones((1, 1)))
              self.assertIsNotNone(lowercase_)


