






import json
import os
from functools import lru_cache
from typing import List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


_UpperCAmelCase					:						Optional[Any]         =     logging.get_logger(__name__)


_UpperCAmelCase					:						Dict         =     {'vocab_file': 'vocab.json', 'merges_file': 'merges.txt'}

# See all MVP models at https://huggingface.co/models?filter=mvp
_UpperCAmelCase					:						List[str]         =     {
    'vocab_file': {
        'RUCAIBox/mvp': 'https://huggingface.co/RUCAIBox/mvp/resolve/main/vocab.json',
    },
    'added_tokens.json': {
        'RUCAIBox/mvp': 'https://huggingface.co/RUCAIBox/mvp/resolve/main/added_tokens.json',
    },
    'merges_file': {
        'RUCAIBox/mvp': 'https://huggingface.co/RUCAIBox/mvp/resolve/main/merges.txt',
    },
}

_UpperCAmelCase					:						Union[str, Any]         =     {
    'RUCAIBox/mvp': 1024,
}






@lru_cache()
def  UpperCAmelCase__						(    ):
   lowerCamelCase									=(
       list(range(ord("""!"""    )			,ord("""~"""    ) + 1    )    ) + list(range(ord("""¡"""    )			,ord("""¬"""    ) + 1    )    ) + list(range(ord("""®"""    )			,ord("""ÿ"""    ) + 1    )    )
   )
   lowerCamelCase									=bs[:]
   lowerCamelCase									=0
   for b in range(2**8    ):
      if b not in bs:
         bs.append(lowerCamelCase    )
         cs.append(2**8 + n    )
         n += 1
   lowerCamelCase									=[chr(lowerCamelCase    ) for n in cs]
   return dict(zip(lowerCamelCase			,lowerCamelCase    )    )






def  UpperCAmelCase__						(    lowerCamelCase	:			List[str]    ):
   lowerCamelCase									=set()
   lowerCamelCase									=word[0]
   for char in word[1:]:
      pairs.add((prev_char, char)    )
      lowerCamelCase									=char
   return pairs







class 		lowerCAmelCase_				(  UpperCamelCase ):
     __lowerCamelCase       :    Any 	=     VOCAB_FILES_NAMES
     __lowerCamelCase       :    int 	=     PRETRAINED_VOCAB_FILES_MAP
     __lowerCamelCase       :    Optional[int] 	=     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __lowerCamelCase       :    List[str] 	=     ['''input_ids''', '''attention_mask''']





     def __init__(	self							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase="replace"							, _UpperCAmelCase="<s>"							, _UpperCAmelCase="</s>"							, _UpperCAmelCase="</s>"							, _UpperCAmelCase="<s>"							, _UpperCAmelCase="<unk>"							, _UpperCAmelCase="<pad>"							, _UpperCAmelCase="<mask>"							, _UpperCAmelCase=False							, **_UpperCAmelCase							, )		->				Any:
        lowerCamelCase									=AddedToken(_UpperCAmelCase							, lstrip=_UpperCAmelCase							, rstrip=_UpperCAmelCase  ) if isinstance(_UpperCAmelCase							, _UpperCAmelCase  ) else bos_token
        lowerCamelCase									=AddedToken(_UpperCAmelCase							, lstrip=_UpperCAmelCase							, rstrip=_UpperCAmelCase  ) if isinstance(_UpperCAmelCase							, _UpperCAmelCase  ) else eos_token
        lowerCamelCase									=AddedToken(_UpperCAmelCase							, lstrip=_UpperCAmelCase							, rstrip=_UpperCAmelCase  ) if isinstance(_UpperCAmelCase							, _UpperCAmelCase  ) else sep_token
        lowerCamelCase									=AddedToken(_UpperCAmelCase							, lstrip=_UpperCAmelCase							, rstrip=_UpperCAmelCase  ) if isinstance(_UpperCAmelCase							, _UpperCAmelCase  ) else cls_token
        lowerCamelCase									=AddedToken(_UpperCAmelCase							, lstrip=_UpperCAmelCase							, rstrip=_UpperCAmelCase  ) if isinstance(_UpperCAmelCase							, _UpperCAmelCase  ) else unk_token
        lowerCamelCase									=AddedToken(_UpperCAmelCase							, lstrip=_UpperCAmelCase							, rstrip=_UpperCAmelCase  ) if isinstance(_UpperCAmelCase							, _UpperCAmelCase  ) else pad_token

        # Mask token behave like a normal word, i.e. include the space before it
        lowerCamelCase									=AddedToken(_UpperCAmelCase							, lstrip=_UpperCAmelCase							, rstrip=_UpperCAmelCase  ) if isinstance(_UpperCAmelCase							, _UpperCAmelCase  ) else mask_token

        super().__init__(
            errors=_UpperCAmelCase							, bos_token=_UpperCAmelCase							, eos_token=_UpperCAmelCase							, unk_token=_UpperCAmelCase							, sep_token=_UpperCAmelCase							, cls_token=_UpperCAmelCase							, pad_token=_UpperCAmelCase							, mask_token=_UpperCAmelCase							, add_prefix_space=_UpperCAmelCase							, **_UpperCAmelCase							, )
        with open(_UpperCAmelCase							, encoding="""utf-8"""  ) as vocab_handle:
           lowerCamelCase									=json.load(_UpperCAmelCase  )
        lowerCamelCase									={v: k for k, v in self.encoder.items()}
        lowerCamelCase									=errors  # how to handle errors in decoding
        lowerCamelCase									=bytes_to_unicode()
        lowerCamelCase									={v: k for k, v in self.byte_encoder.items()}
        with open(_UpperCAmelCase							, encoding="""utf-8"""  ) as merges_handle:
           lowerCamelCase									=merges_handle.read().split("""\n"""  )[1:-1]
        lowerCamelCase									=[tuple(merge.split()  ) for merge in bpe_merges]
        lowerCamelCase									=dict(zip(_UpperCAmelCase							, range(len(_UpperCAmelCase  )  )  )  )
        lowerCamelCase									={}
        lowerCamelCase									=add_prefix_space

        # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
        lowerCamelCase									=re.compile(R"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""  )





     @property
     def 						lowerCamelCase							(	self  )		->				Any:
        return len(self.encoder  )





     def 						lowerCamelCase							(	self  )		->				List[str]:
        return dict(self.encoder							, **self.added_tokens_encoder  )





     def 						lowerCamelCase							(	self							, _UpperCAmelCase  )		->				List[Any]:
        if token in self.cache:
           return self.cache[token]
        lowerCamelCase									=tuple(_UpperCAmelCase  )
        lowerCamelCase									=get_pairs(_UpperCAmelCase  )

        if not pairs:
           return token

        while True:
           lowerCamelCase									=min(_UpperCAmelCase							, key=lambda _UpperCAmelCase  : self.bpe_ranks.get(_UpperCAmelCase							, float("""inf"""  )  )  )
           if bigram not in self.bpe_ranks:
              break
           lowerCamelCase				,       lowerCamelCase									=bigram
           lowerCamelCase									=[]
           lowerCamelCase									=0
           while i < len(_UpperCAmelCase  ):
              try:
                 lowerCamelCase									=word.index(_UpperCAmelCase							, _UpperCAmelCase  )
              except ValueError:
                 new_word.extend(word[i:]  )
                 break
              else:
                 new_word.extend(word[i:j]  )
                 lowerCamelCase									=j

              if word[i] == first and i < len(_UpperCAmelCase  ) - 1 and word[i + 1] == second:
                 new_word.append(first + second  )
                 i += 2
              else:
                 new_word.append(word[i]  )
                 i += 1
           lowerCamelCase									=tuple(_UpperCAmelCase  )
           lowerCamelCase									=new_word
           if len(_UpperCAmelCase  ) == 1:
              break
           else:
              lowerCamelCase									=get_pairs(_UpperCAmelCase  )
        lowerCamelCase									=""" """.join(_UpperCAmelCase  )
        lowerCamelCase									=word
        return word





     def 						lowerCamelCase							(	self							, _UpperCAmelCase  )		->				Tuple:
        lowerCamelCase									=[]
        for token in re.findall(self.pat							, _UpperCAmelCase  ):
           lowerCamelCase									="""""".join(
               self.byte_encoder[b] for b in token.encode("""utf-8"""  )  )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
           bpe_tokens.extend(bpe_token for bpe_token in self.bpe(_UpperCAmelCase  ).split(""" """  )  )
        return bpe_tokens





     def 						lowerCamelCase							(	self							, _UpperCAmelCase  )		->				Tuple:
        return self.encoder.get(_UpperCAmelCase							, self.encoder.get(self.unk_token  )  )





     def 						lowerCamelCase							(	self							, _UpperCAmelCase  )		->				Any:
        return self.decoder.get(_UpperCAmelCase  )





     def 						lowerCamelCase							(	self							, _UpperCAmelCase  )		->				Tuple:
        lowerCamelCase									="""""".join(_UpperCAmelCase  )
        lowerCamelCase									=bytearray([self.byte_decoder[c] for c in text]  ).decode("""utf-8"""							, errors=self.errors  )
        return text





     def 						lowerCamelCase							(	self							, _UpperCAmelCase							, _UpperCAmelCase = None  )		->				Tuple[str]:
        if not os.path.isdir(_UpperCAmelCase  ):
           logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''  )
           return
        lowerCamelCase									=os.path.join(
            _UpperCAmelCase							, (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]  )
        lowerCamelCase									=os.path.join(
            _UpperCAmelCase							, (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""merges_file"""]  )

        with open(_UpperCAmelCase							, """w"""							, encoding="""utf-8"""  ) as f:
           f.write(json.dumps(self.encoder							, indent=2							, sort_keys=_UpperCAmelCase							, ensure_ascii=_UpperCAmelCase  ) + """\n"""  )

        lowerCamelCase									=0
        with open(_UpperCAmelCase							, """w"""							, encoding="""utf-8"""  ) as writer:
           writer.write("""#version: 0.2\n"""  )
           for bpe_tokens, token_index in sorted(self.bpe_ranks.items()							, key=lambda _UpperCAmelCase  : kv[1]  ):
              if index != token_index:
                 logger.warning(
                     F'''Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'''
                     """ Please check that the tokenizer is not corrupted!"""  )
                 lowerCamelCase									=token_index
              writer.write(""" """.join(_UpperCAmelCase  ) + """\n"""  )
              index += 1

        return vocab_file, merge_file





     def 						lowerCamelCase							(	self							, _UpperCAmelCase							, _UpperCAmelCase = None  )		->				List[int]:
        if token_ids_a is None:
           return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
        lowerCamelCase									=[self.cls_token_id]
        lowerCamelCase									=[self.sep_token_id]
        return cls + token_ids_a + sep + sep + token_ids_a + sep





     def 						lowerCamelCase							(	self							, _UpperCAmelCase							, _UpperCAmelCase = None							, _UpperCAmelCase = False  )		->				List[int]:
        if already_has_special_tokens:
           return super().get_special_tokens_mask(
               token_ids_a=_UpperCAmelCase							, token_ids_a=_UpperCAmelCase							, already_has_special_tokens=_UpperCAmelCase  )

        if token_ids_a is None:
           return [1] + ([0] * len(_UpperCAmelCase  )) + [1]
        return [1] + ([0] * len(_UpperCAmelCase  )) + [1, 1] + ([0] * len(_UpperCAmelCase  )) + [1]





     def 						lowerCamelCase							(	self							, _UpperCAmelCase							, _UpperCAmelCase = None  )		->				List[int]:
        lowerCamelCase									=[self.sep_token_id]
        lowerCamelCase									=[self.cls_token_id]

        if token_ids_a is None:
           return len(cls + token_ids_a + sep  ) * [0]
        return len(cls + token_ids_a + sep + sep + token_ids_a + sep  ) * [0]





     def 						lowerCamelCase							(	self							, _UpperCAmelCase							, _UpperCAmelCase=False							, **_UpperCAmelCase  )		->				Dict:
        lowerCamelCase									=kwargs.pop("""add_prefix_space"""							, self.add_prefix_space  )
        if (is_split_into_words or add_prefix_space) and (len(_UpperCAmelCase  ) > 0 and not text[0].isspace()):
           lowerCamelCase									=""" """ + text
        return (text, kwargs)







import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
   import jax
   import jax.numpy as jnp
   import numpy as np

   from transformers.modeling_flax_pytorch_utils import (
       convert_pytorch_state_dict_to_flax,
       load_flax_weights_in_pytorch_model,
   )
   from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
   import torch


@require_flax
class 		lowerCAmelCase_				:





     def __init__(	self							, _UpperCAmelCase							, _UpperCAmelCase=14							, _UpperCAmelCase=7							, _UpperCAmelCase=True							, _UpperCAmelCase=True							, _UpperCAmelCase=True							, _UpperCAmelCase=99							, _UpperCAmelCase=32							, _UpperCAmelCase=5							, _UpperCAmelCase=4							, _UpperCAmelCase=37							, _UpperCAmelCase="gelu"							, _UpperCAmelCase=0.1							, _UpperCAmelCase=0.1							, _UpperCAmelCase=512							, _UpperCAmelCase=0.02							, _UpperCAmelCase=None							, )		->				str:
        lowerCamelCase									=parent
        lowerCamelCase									=batch_size
        lowerCamelCase									=seq_length
        lowerCamelCase									=is_training
        lowerCamelCase									=use_input_mask
        lowerCamelCase									=use_labels
        lowerCamelCase									=vocab_size
        lowerCamelCase									=d_model
        lowerCamelCase									=num_hidden_layers
        lowerCamelCase									=num_attention_heads
        lowerCamelCase									=ffn_dim
        lowerCamelCase									=activation_function
        lowerCamelCase									=activation_dropout
        lowerCamelCase									=attention_dropout
        lowerCamelCase									=max_position_embeddings
        lowerCamelCase									=initializer_range
        lowerCamelCase									=None
        lowerCamelCase									=0
        lowerCamelCase									=2
        lowerCamelCase									=1





     def 						lowerCamelCase							(	self  )		->				List[str]:
        lowerCamelCase									=np.clip(ids_tensor([self.batch_size, self.seq_length]							, self.vocab_size  )							, 3							, self.vocab_size  )

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

        lowerCamelCase									=XGLMConfig(
            vocab_size=self.vocab_size							, d_model=self.hidden_size							, num_layers=self.num_hidden_layers							, attention_heads=self.num_attention_heads							, ffn_dim=self.ffn_dim							, activation_function=self.activation_function							, activation_dropout=self.activation_dropout							, attention_dropout=self.attention_dropout							, max_position_embeddings=self.max_position_embeddings							, initializer_range=self.initializer_range							, use_cache=_UpperCAmelCase							, bos_token_id=self.bos_token_id							, eos_token_id=self.eos_token_id							, pad_token_id=self.pad_token_id							, )

        return (config, input_ids, input_mask)





     def 						lowerCamelCase							(	self  )		->				Tuple:
        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 						lowerCamelCase							(	self  )		->				int:
        lowerCamelCase				,       lowerCamelCase				,       lowerCamelCase									=self.prepare_config_and_inputs()

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

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





     def 						lowerCamelCase							(	self							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )		->				Optional[Any]:
        lowerCamelCase									=20
        lowerCamelCase									=model_class_name(_UpperCAmelCase  )

        lowerCamelCase									=model.init_cache(input_ids.shape[0]							, _UpperCAmelCase  )
        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=_UpperCAmelCase							, past_key_values=_UpperCAmelCase							, position_ids=_UpperCAmelCase							, )

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

        lowerCamelCase									=model(_UpperCAmelCase  )

        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 						lowerCamelCase							(	self							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )		->				List[Any]:
        lowerCamelCase									=20
        lowerCamelCase									=model_class_name(_UpperCAmelCase  )

        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]							, _UpperCAmelCase  )
        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=_UpperCAmelCase							, past_key_values=_UpperCAmelCase							, position_ids=_UpperCAmelCase							, )
        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=_UpperCAmelCase							, position_ids=_UpperCAmelCase							, )

        lowerCamelCase									=model(_UpperCAmelCase							, attention_mask=_UpperCAmelCase  )
        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_sentencepiece
@require_flax
class 		lowerCAmelCase_				(  UpperCamelCase				,     UpperCamelCase				,     unittest.TestCase ):
     __lowerCamelCase       :    Tuple 	=     (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
     __lowerCamelCase       :    Dict 	=     (FlaxXGLMForCausalLM,) if is_flax_available() else ()





     def 						lowerCamelCase							(	self  )		->				List[str]:
        lowerCamelCase									=FlaxXGLMModelTester(self  )





     def 						lowerCamelCase							(	self  )		->				List[str]:
        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(_UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )





     def 						lowerCamelCase							(	self  )		->				Tuple:
        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(
               _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )





     @slow
     def 						lowerCamelCase							(	self  )		->				Tuple:
        lowerCamelCase									=XGLMTokenizer.from_pretrained("""XGLM"""							, padding_side="""left"""  )
        lowerCamelCase									=tokenizer(["""Hello this is a long string""", """Hey"""]							, return_tensors="""np"""							, padding=_UpperCAmelCase							, truncation=_UpperCAmelCase  )

        lowerCamelCase									=FlaxXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""  )
        lowerCamelCase									=1
        lowerCamelCase									=False

        lowerCamelCase									=jax.jit(model.generate  )

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

        lowerCamelCase									=tokenizer.batch_decode(_UpperCAmelCase							, skip_special_tokens=_UpperCAmelCase  )

        lowerCamelCase									=[
            """Hello this is a long string of questions, but I'm not sure if I'm""",
            """Hey, I'm a newbie to the forum and I'""",
        ]

        self.assertListEqual(_UpperCAmelCase							, _UpperCAmelCase  )





     @is_pt_flax_cross_test
     def 						lowerCamelCase							(	self  )		->				Dict:
        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(_UpperCAmelCase							, _UpperCAmelCase  )
              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(_UpperCAmelCase							, _UpperCAmelCase  )

              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(_UpperCAmelCase  ):
                 lowerCamelCase									=0
                 lowerCamelCase									=1
                 lowerCamelCase									=0
                 lowerCamelCase									=1
              lowerCamelCase									=pt_model_class(_UpperCAmelCase  ).eval()
              # Flax models don't use the `use_cache` option and cache is not returned as a default.
              # So we disable `use_cache` here for PyTorch model.
              lowerCamelCase									=False
              lowerCamelCase									=model_class(_UpperCAmelCase							, dtype=jnp.floataa  )
              lowerCamelCase									=convert_pytorch_state_dict_to_flax(pt_model.state_dict()							, _UpperCAmelCase  )
              lowerCamelCase									=fx_state

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

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

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

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





     @is_pt_flax_cross_test
     def 						lowerCamelCase							(	self  )		->				Optional[Any]:
        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(_UpperCAmelCase							, _UpperCAmelCase  )
              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(_UpperCAmelCase							, _UpperCAmelCase  )

              lowerCamelCase									=pt_model_class(_UpperCAmelCase  ).eval()
              lowerCamelCase									=False
              lowerCamelCase									=model_class(_UpperCAmelCase							, dtype=jnp.floataa  )

              lowerCamelCase									=load_flax_weights_in_pytorch_model(_UpperCAmelCase							, 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(_UpperCAmelCase  ):
                 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(**_UpperCAmelCase  ).to_tuple()

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

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

              with torch.no_grad():
                 lowerCamelCase									=pt_model_loaded(**_UpperCAmelCase  ).to_tuple()

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





     @slow
     def 						lowerCamelCase							(	self  )		->				int:
        for model_class_name in self.all_model_classes:
           lowerCamelCase									=model_class_name.from_pretrained("""facebook/xglm-564M"""  )
           lowerCamelCase									=model(np.ones((1, 1)  )  )
           self.assertIsNotNone(_UpperCAmelCase  )