


import copy
import tempfile
import unittest

from transformers import MaMaaaConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device
from transformers.utils import cached_property

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import MaMaaaForConditionalGeneration, MaMaaaModel, MaMaaaTokenizer
 from transformers.models.mam_aaa.modeling_mam_aaa import MaMaaaDecoder, MaMaaaEncoder
def   _UpperCamelCase				(						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None  ,						)						->Optional[Any]:
   if attention_mask is None:
      UpperCAmelCase						=     input_ids.ne(config.pad_token_id   )
   if decoder_attention_mask is None:
      UpperCAmelCase						=     decoder_input_ids.ne(config.pad_token_id   )
   if head_mask is None:
      UpperCAmelCase						=     torch.ones(config.encoder_layers  ,						config.encoder_attention_heads  ,						device=__UpperCamelCase   )
   if decoder_head_mask is None:
      UpperCAmelCase						=     torch.ones(config.decoder_layers  ,						config.decoder_attention_heads  ,						device=__UpperCamelCase   )
   if cross_attn_head_mask is None:
      UpperCAmelCase						=     torch.ones(config.decoder_layers  ,						config.decoder_attention_heads  ,						device=__UpperCamelCase   )
   return {
       "input_ids": input_ids,
       "decoder_input_ids": decoder_input_ids,
       "attention_mask": attention_mask,
       "decoder_attention_mask": attention_mask,
       "head_mask": head_mask,
       "decoder_head_mask": decoder_head_mask,
       "cross_attn_head_mask": cross_attn_head_mask,
   }






class   __lowercase				:





      def __init__(		self						:					Dict					,       __lowerCamelCase						:					Any					,       __lowerCamelCase						:					Optional[Any]=1_3					,       __lowerCamelCase						:					Union[str, Any]=7					,       __lowerCamelCase						:					int=True					,       __lowerCamelCase						:					Dict=False					,       __lowerCamelCase						:					Optional[int]=9_9					,       __lowerCamelCase						:					Any=1_6					,       __lowerCamelCase						:					Optional[int]=2					,       __lowerCamelCase						:					Optional[int]=4					,       __lowerCamelCase						:					Any=4					,       __lowerCamelCase						:					List[Any]="relu"					,       __lowerCamelCase						:					List[Any]=0.1					,       __lowerCamelCase						:					Optional[int]=0.1					,       __lowerCamelCase						:					Any=0.0					,       __lowerCamelCase						:					List[Any]=0.0					,       __lowerCamelCase						:					Dict=2_0					,       __lowerCamelCase						:					Optional[Any]=2					,       __lowerCamelCase						:					Optional[Any]=1					,       __lowerCamelCase						:					Optional[int]=0					,       )    ->   List[Any]:






         """simple docstring"""







         UpperCAmelCase						=     parent
         UpperCAmelCase						=     batch_size
         UpperCAmelCase						=     seq_length
         UpperCAmelCase						=     is_training
         UpperCAmelCase						=     use_labels
         UpperCAmelCase						=     vocab_size
         UpperCAmelCase						=     hidden_size
         UpperCAmelCase						=     num_hidden_layers
         UpperCAmelCase						=     num_attention_heads
         UpperCAmelCase						=     intermediate_size
         UpperCAmelCase						=     hidden_act
         UpperCAmelCase						=     hidden_dropout_prob
         UpperCAmelCase						=     attention_probs_dropout_prob
         UpperCAmelCase						=     encoder_layerdrop
         UpperCAmelCase						=     decoder_layerdrop
         UpperCAmelCase						=     max_position_embeddings
         UpperCAmelCase						=     eos_token_id
         UpperCAmelCase						=     pad_token_id
         UpperCAmelCase						=     bos_token_id





      def        _lowercase				(		self						:					str       )    ->   Dict:






         """simple docstring"""







         UpperCAmelCase						=     ids_tensor([self.batch_size, self.seq_length]					,       self.vocab_size       )
         UpperCAmelCase						=     self.eos_token_id  # Eos Token
         UpperCAmelCase						=     ids_tensor([self.batch_size, self.seq_length]					,       self.vocab_size       )

         # we need to clamp the input ids here to avoid having pad token in between
         # this is because for M2M100 the position_ids are prepared such that
         # all pad tokens have pos id = 2 and rest are between 2..seq_length
         # and the seq_length here is seq_length - num_pad_tokens
         # but when using past, there is no way of knowing if the past input ids had
         # pad tokens in them, which results in incorrect seq_lenth and which in turn results in
         # position_ids being off by num_pad_tokens in past input
         UpperCAmelCase						=     input_ids.clamp(self.pad_token_id + 1       )
         UpperCAmelCase						=     decoder_input_ids.clamp(self.pad_token_id + 1       )

         UpperCAmelCase						=     self.get_config()
         UpperCAmelCase						=     prepare_mam_aaa_inputs_dict(_SCREAMING_SNAKE_CASE					,       _SCREAMING_SNAKE_CASE					,       _SCREAMING_SNAKE_CASE       )
         return config, inputs_dict





      def        _lowercase				(		self						:					Dict       )    ->   Union[str, Any]:






         """simple docstring"""







         return MaMaaaConfig(
             vocab_size=self.vocab_size					,       d_model=self.hidden_size					,       encoder_layers=self.num_hidden_layers					,       decoder_layers=self.num_hidden_layers					,       encoder_attention_heads=self.num_attention_heads					,       decoder_attention_heads=self.num_attention_heads					,       encoder_ffn_dim=self.intermediate_size					,       decoder_ffn_dim=self.intermediate_size					,       dropout=self.hidden_dropout_prob					,       attention_dropout=self.attention_probs_dropout_prob					,       encoder_layerdrop=self.encoder_layerdrop					,       decoder_layerdrop=self.decoder_layerdrop					,       max_position_embeddings=self.max_position_embeddings					,       eos_token_id=self.eos_token_id					,       bos_token_id=self.bos_token_id					,       pad_token_id=self.pad_token_id					,       )





      def        _lowercase				(		self						:					List[Any]       )    ->   Any:






         """simple docstring"""







         UpperCAmelCase				,  UpperCAmelCase						=     self.prepare_config_and_inputs()
         return config, inputs_dict





      def        _lowercase				(		self						:					Any					,       __lowerCamelCase						:					Optional[int]					,       __lowerCamelCase						:					Optional[Any]       )    ->   int:






         """simple docstring"""







         UpperCAmelCase						=     MaMaaaModel(config=_SCREAMING_SNAKE_CASE       ).get_decoder().to(_SCREAMING_SNAKE_CASE       ).eval()
         UpperCAmelCase						=     inputs_dict["""input_ids"""]
         UpperCAmelCase						=     inputs_dict["""attention_mask"""]
         UpperCAmelCase						=     inputs_dict["""head_mask"""]

         # first forward pass
         UpperCAmelCase						=     model(_SCREAMING_SNAKE_CASE					,       attention_mask=_SCREAMING_SNAKE_CASE					,       head_mask=_SCREAMING_SNAKE_CASE					,       use_cache=_SCREAMING_SNAKE_CASE       )

         UpperCAmelCase				,  UpperCAmelCase						=     outputs.to_tuple()

         # create hypothetical multiple next token and extent to next_input_ids
         UpperCAmelCase						=     ids_tensor((self.batch_size, 3)					,       config.vocab_size       )
         UpperCAmelCase						=     ids_tensor((self.batch_size, 3)					,       2       )

         # append to next input_ids and
         UpperCAmelCase						=     torch.cat([input_ids, next_tokens]					,       dim=-1       )
         UpperCAmelCase						=     torch.cat([attention_mask, next_attn_mask]					,       dim=-1       )

         UpperCAmelCase						=     model(_SCREAMING_SNAKE_CASE					,       attention_mask=_SCREAMING_SNAKE_CASE       )["""last_hidden_state"""]
         UpperCAmelCase						=     model(_SCREAMING_SNAKE_CASE					,       attention_mask=_SCREAMING_SNAKE_CASE					,       past_key_values=_SCREAMING_SNAKE_CASE       )[
             """last_hidden_state"""
         ]

         # select random slice
         UpperCAmelCase						=     ids_tensor((1,)					,       output_from_past.shape[-1]       ).item()
         UpperCAmelCase						=     output_from_no_past[:, -3:, random_slice_idx].detach()
         UpperCAmelCase						=     output_from_past[:, :, random_slice_idx].detach()

         self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]       )

         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(_SCREAMING_SNAKE_CASE					,       _SCREAMING_SNAKE_CASE					,       atol=1e-2       )       )





      def        _lowercase				(		self						:					Tuple					,       __lowerCamelCase						:					Optional[Any]					,       __lowerCamelCase						:					Dict       )    ->   Optional[int]:






         """simple docstring"""







         UpperCAmelCase						=     MaMaaaModel(config=_SCREAMING_SNAKE_CASE       ).to(_SCREAMING_SNAKE_CASE       ).eval()
         UpperCAmelCase						=     model(**_SCREAMING_SNAKE_CASE       )

         UpperCAmelCase						=     outputs.encoder_last_hidden_state
         UpperCAmelCase						=     outputs.last_hidden_state

         with tempfile.TemporaryDirectory() as tmpdirname:
            UpperCAmelCase						=     model.get_encoder()
            encoder.save_pretrained(_SCREAMING_SNAKE_CASE       )
            UpperCAmelCase						=     MaMaaaEncoder.from_pretrained(_SCREAMING_SNAKE_CASE       ).to(_SCREAMING_SNAKE_CASE       )

         UpperCAmelCase						=     encoder(inputs_dict["""input_ids"""]					,       attention_mask=inputs_dict["""attention_mask"""]       )[
             0
         ]

         self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1e-3       )

         with tempfile.TemporaryDirectory() as tmpdirname:
            UpperCAmelCase						=     model.get_decoder()
            decoder.save_pretrained(_SCREAMING_SNAKE_CASE       )
            UpperCAmelCase						=     MaMaaaDecoder.from_pretrained(_SCREAMING_SNAKE_CASE       ).to(_SCREAMING_SNAKE_CASE       )

         UpperCAmelCase						=     decoder(
             input_ids=inputs_dict["""decoder_input_ids"""]					,       attention_mask=inputs_dict["""decoder_attention_mask"""]					,       encoder_hidden_states=_SCREAMING_SNAKE_CASE					,       encoder_attention_mask=inputs_dict["""attention_mask"""]					,       )[0]

         self.parent.assertTrue((last_hidden_state_a - last_hidden_state).abs().max().item() < 1e-3       )






@require_torch
class   __lowercase				(  __snake_case ,					__snake_case ,					__snake_case ,					unittest.TestCase  ):
      UpperCamelCase								=       (
          (
              MaMaaaModel,
              MaMaaaForConditionalGeneration,
          )
          if is_torch_available()
          else ()
      )
      UpperCamelCase								=       (MaMaaaForConditionalGeneration,) if is_torch_available() else ()
      UpperCamelCase								=       (
          {
              '''conversational''': MaMaaaForConditionalGeneration,
              '''feature-extraction''': MaMaaaModel,
              '''summarization''': MaMaaaForConditionalGeneration,
              '''text2text-generation''': MaMaaaForConditionalGeneration,
              '''translation''': MaMaaaForConditionalGeneration,
          }
          if is_torch_available()
          else {}
      )
      UpperCamelCase								=       True
      UpperCamelCase								=       True
      UpperCamelCase								=       False
      UpperCamelCase								=       False





      def        _lowercase				(		self						:					List[Any]					,       __lowerCamelCase						:					Union[str, Any]					,       __lowerCamelCase						:					Optional[int]					,       __lowerCamelCase						:					List[Any]					,       __lowerCamelCase						:					Dict					,       __lowerCamelCase						:					Tuple       )    ->   Union[str, Any]:






         """simple docstring"""







         if pipeline_test_casse_name == "TranslationPipelineTests":
            # Get `ValueError: Translation requires a `src_lang` and a `tgt_lang` for this model`.
            # `M2M100Config` was never used in pipeline tests: cannot create a simple tokenizer.
            return True

         return False





      def        _lowercase				(		self						:					Optional[int]       )    ->   Optional[int]:






         """simple docstring"""







         UpperCAmelCase						=     MaMaaaModelTester(self       )
         UpperCAmelCase						=     ConfigTester(self					,       config_class=_SCREAMING_SNAKE_CASE       )





      def        _lowercase				(		self						:					List[Any]       )    ->   Optional[Any]:






         """simple docstring"""







         self.config_tester.run_common_tests()





      def        _lowercase				(		self						:					List[str]       )    ->   Dict:






         """simple docstring"""







         UpperCAmelCase				,  UpperCAmelCase						=     self.model_tester.prepare_config_and_inputs()
         for model_class in self.all_model_classes:
            UpperCAmelCase						=     model_class(_SCREAMING_SNAKE_CASE       )

            with tempfile.TemporaryDirectory() as tmpdirname:
               model.save_pretrained(_SCREAMING_SNAKE_CASE       )
               UpperCAmelCase				,  UpperCAmelCase						=     model_class.from_pretrained(_SCREAMING_SNAKE_CASE					,       output_loading_info=_SCREAMING_SNAKE_CASE       )
            self.assertEqual(info["""missing_keys"""]					,       []       )





      def        _lowercase				(		self						:					Union[str, Any]       )    ->   Optional[int]:






         """simple docstring"""







         UpperCAmelCase						=     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_decoder_model_past_large_inputs(*_SCREAMING_SNAKE_CASE       )





      def        _lowercase				(		self						:					Any       )    ->   Dict:






         """simple docstring"""







         UpperCAmelCase						=     self.model_tester.prepare_config_and_inputs_for_common()
         self.model_tester.check_encoder_decoder_model_standalone(*_SCREAMING_SNAKE_CASE       )





      def        _lowercase				(		self						:					Tuple       )    ->   Dict:






         """simple docstring"""







         UpperCAmelCase				,  UpperCAmelCase						=     self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in (MaMaaaModel, MaMaaaForConditionalGeneration):
            UpperCAmelCase						=     model_class(_SCREAMING_SNAKE_CASE       )
            model.to(_SCREAMING_SNAKE_CASE       )
            model.eval()

            UpperCAmelCase						=     copy.deepcopy(self._prepare_for_class(_SCREAMING_SNAKE_CASE					,       _SCREAMING_SNAKE_CASE       )       )

            if not self.is_encoder_decoder:
               UpperCAmelCase						=     inputs["""input_ids"""]
               del inputs["input_ids"]
            else:
               UpperCAmelCase						=     inputs["""input_ids"""]
               UpperCAmelCase						=     inputs.get("""decoder_input_ids"""					,       _SCREAMING_SNAKE_CASE       )
               del inputs["input_ids"]
               inputs.pop("""decoder_input_ids"""					,       _SCREAMING_SNAKE_CASE       )

            UpperCAmelCase						=     model.get_input_embeddings()
            if not self.is_encoder_decoder:
               UpperCAmelCase						=     wte(_SCREAMING_SNAKE_CASE       )
            else:
               UpperCAmelCase						=     wte(_SCREAMING_SNAKE_CASE       )
               UpperCAmelCase						=     wte(_SCREAMING_SNAKE_CASE       )

            with torch.no_grad():
               model(**_SCREAMING_SNAKE_CASE       )[0]





      def        _lowercase				(		self						:					int       )    ->   int:






         """simple docstring"""







         UpperCAmelCase				,  UpperCAmelCase						=     self.model_tester.prepare_config_and_inputs()
         UpperCAmelCase						=     input_dict["""input_ids"""]
         UpperCAmelCase						=     input_ids.ne(1       ).to(_SCREAMING_SNAKE_CASE       )
         UpperCAmelCase						=     MaMaaaForConditionalGeneration(_SCREAMING_SNAKE_CASE       ).eval().to(_SCREAMING_SNAKE_CASE       )
         if torch_device == "cuda":
            model.half()
         model.generate(_SCREAMING_SNAKE_CASE					,       attention_mask=_SCREAMING_SNAKE_CASE       )
         model.generate(num_beams=4					,       do_sample=_SCREAMING_SNAKE_CASE					,       early_stopping=_SCREAMING_SNAKE_CASE					,       num_return_sequences=3       )

def   _UpperCamelCase				(						lowerCAmelCase_   )						->Any:
   return torch.tensor(__UpperCamelCase  ,						dtype=torch.long  ,						device=__UpperCamelCase   )


__a     =			1E-4







@require_torch
@require_sentencepiece
@require_tokenizers
@slow
class   __lowercase				(  unittest.TestCase  ):





      @cached_property
      def        _lowercase				(		self						:					Any       )    ->   Any:






         """simple docstring"""







         return MaMaaaTokenizer.from_pretrained("""facebook/m2m100_418M"""       )





      def        _lowercase				(		self						:					Tuple       )    ->   str:






         """simple docstring"""







         UpperCAmelCase						=     MaMaaaModel.from_pretrained("""facebook/m2m100_418M"""       ).to(_SCREAMING_SNAKE_CASE       )
         UpperCAmelCase						=     _long_tensor([[1_2_8_0_2_8, 9_8, 1_2, 3_0_5_2_7, 2_7_3_2, 1_5_9, 7_7_5_5, 6_1_9_0_4, 3_9_1_4_4, 3_8, 2]]       )
         UpperCAmelCase						=     _long_tensor([[2, 1_2_8_0_2_8, 9_8, 1_2, 3_0_5_2_7, 2_7_3_2, 1_5_9, 7_7_5_5, 6_1_9_0_4, 3_9_1_4_4, 3_8]]       )
         UpperCAmelCase						=     prepare_mam_aaa_inputs_dict(model.config					,       _SCREAMING_SNAKE_CASE					,       _SCREAMING_SNAKE_CASE       )
         with torch.no_grad():
            UpperCAmelCase						=     model(**_SCREAMING_SNAKE_CASE       )[0]
         UpperCAmelCase						=     torch.Size((1, 1_1, 1_0_2_4)       )
         self.assertEqual(output.shape					,       _SCREAMING_SNAKE_CASE       )
         # change to expected output here
         UpperCAmelCase						=     torch.tensor(
             [[-0.7_780, -0.1_676, 0.1_038], [-6.7_556, -1.3_992, 0.0_567], [-7.5_383, -0.5_920, -0.2_779]]					,       device=_SCREAMING_SNAKE_CASE       )
         self.assertTrue(torch.allclose(output[:, :3, :3]					,       _SCREAMING_SNAKE_CASE					,       atol=_SCREAMING_SNAKE_CASE       )       )





      def        _lowercase				(		self						:					str       )    ->   Tuple:






         """simple docstring"""







         UpperCAmelCase						=     MaMaaaForConditionalGeneration.from_pretrained("""facebook/m2m100_418M"""       ).to(_SCREAMING_SNAKE_CASE       )

         # change to intended input
         UpperCAmelCase						=     _long_tensor([[1_2_8_0_2_8, 9_8, 1_2, 3_0_5_2_7, 2_7_3_2, 1_5_9, 7_7_5_5, 6_1_9_0_4, 3_9_1_4_4, 3_8, 2]]       )
         UpperCAmelCase						=     _long_tensor([[2, 1_2_8_0_2_8, 9_8, 1_2, 3_0_5_2_7, 2_7_3_2, 1_5_9, 7_7_5_5, 6_1_9_0_4, 3_9_1_4_4, 3_8]]       )
         UpperCAmelCase						=     prepare_mam_aaa_inputs_dict(model.config					,       _SCREAMING_SNAKE_CASE					,       _SCREAMING_SNAKE_CASE       )
         with torch.no_grad():
            UpperCAmelCase						=     model(**_SCREAMING_SNAKE_CASE       )[0]
         UpperCAmelCase						=     torch.Size((1, 1_1, model.config.vocab_size)       )
         self.assertEqual(output.shape					,       _SCREAMING_SNAKE_CASE       )
         # change to expected output here
         UpperCAmelCase						=     torch.tensor(
             [[-1.0_448, -1.0_411, 3.7_992], [-3.2_191, -3.2_386, -1.3_451], [-3.6_210, -3.5_993, 0.4_925]]					,       device=_SCREAMING_SNAKE_CASE       )
         self.assertTrue(torch.allclose(output[:, :3, :3]					,       _SCREAMING_SNAKE_CASE					,       atol=_SCREAMING_SNAKE_CASE       )       )





      def        _lowercase				(		self						:					Dict       )    ->   List[Any]:






         """simple docstring"""







         UpperCAmelCase						=     MaMaaaForConditionalGeneration.from_pretrained("""facebook/m2m100_418M"""       ).to(_SCREAMING_SNAKE_CASE       )
         UpperCAmelCase						=     MaMaaaTokenizer.from_pretrained("""facebook/m2m100_418M"""					,       src_lang="""fr"""					,       tgt_lang="""en"""       )

         UpperCAmelCase						=     [
             """L'affaire NSA souligne l'absence totale de débat sur le renseignement""",
             """Selon moi, il y a deux niveaux de réponse de la part du gouvernement français.""",
             """Lorsque François Hollande téléphone à Barack Obama ou quand le ministre des affaires étrangères Laurent"""
             """ Fabius convoque l'ambassadeur des Etats-Unis, ils réagissent à une vraie découverte, qui est celle de"""
             """ l'ampleur de la surveillance américaine sur l'ensemble des communications en France.""",
         ]

         # The below article tests that we don't add any hypotheses outside of the top n_beams
         UpperCAmelCase						=     tokenizer(_SCREAMING_SNAKE_CASE					,       padding=_SCREAMING_SNAKE_CASE					,       return_tensors="""pt"""       )

         UpperCAmelCase						=     model.generate(
             input_ids=dct["""input_ids"""].to(_SCREAMING_SNAKE_CASE       )					,       attention_mask=dct["""attention_mask"""].to(_SCREAMING_SNAKE_CASE       )					,       num_beams=5					,       forced_bos_token_id=tokenizer.get_lang_id("""en"""       )					,       )

         UpperCAmelCase						=     [
             """The NSA case highlights the total absence of intelligence debate""",
             """I think there are two levels of response from the French government.""",
             """When François Hollande calls Barack Obama or when Foreign Minister Laurent Fabius calls the U.S."""
             """ Ambassador, they respond to a real discovery, which is that of the scale of U.S. surveillance on all"""
             """ communications in France.""",
         ]

         UpperCAmelCase						=     tokenizer.batch_decode(
             hypotheses_batch.tolist()					,       clean_up_tokenization_spaces=_SCREAMING_SNAKE_CASE					,       skip_special_tokens=_SCREAMING_SNAKE_CASE       )
         assert generated == expected_en







from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


__a     =			logging.get_logger(__name__)
def   _UpperCamelCase				(						lowerCAmelCase_  ,						lowerCAmelCase_   )						->Tuple:
   try:
      with open(lowerCAmelCase_  ,						"""rb"""   ) as flax_state_f:
         UpperCAmelCase							=     from_bytes(lowerCAmelCase_  ,						flax_state_f.read()   )
   except UnpicklingError as e:
      try:
         with open(lowerCAmelCase_   ) as f:
            if f.read().startswith("""version"""   ):
               raise OSError(
                   """You seem to have cloned a repository without having git-lfs installed. Please"""
                   """ install git-lfs and run `git lfs install` followed by `git lfs pull` in the"""
                   """ folder you cloned."""   )
            else:
               raise ValueError from e
      except (UnicodeDecodeError, ValueError):
         raise EnvironmentError(F"""Unable to convert {model_file} to Flax deserializable object. """   )

   return load_flax_weights_in_pytorch_model(lowerCAmelCase_  ,						lowerCAmelCase_   )







def   _UpperCamelCase				(						lowerCAmelCase_  ,						lowerCAmelCase_   )						->Dict:

   try:
      import torch  # noqa: F401
   except ImportError:
      logger.error(
          """Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see"""
          """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
          """ instructions."""   )
      raise

   # check if we have bf16 weights
   UpperCAmelCase							=     flatten_dict(jax.tree_util.tree_map(lambda lowerCAmelCase_   : x.dtype == jnp.bfloataa  ,						lowerCAmelCase_   )   ).values()
   if any(lowerCAmelCase_   ):
      # convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

      # and bf16 is not fully supported in PT yet.
      logger.warning(
          """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
          """before loading those in PyTorch model."""   )
      UpperCAmelCase							=     jax.tree_util.tree_map(
          lambda lowerCAmelCase_   : params.astype(np.floataa   ) if params.dtype == jnp.bfloataa else params  ,						lowerCAmelCase_   )

   UpperCAmelCase							=     """"""

   UpperCAmelCase							=     flatten_dict(lowerCAmelCase_  ,						sep="""."""   )
   UpperCAmelCase							=     pt_model.state_dict()

   # keep track of unexpected & missing keys
   UpperCAmelCase							=     []
   UpperCAmelCase							=     set(pt_model_dict.keys()   )

   for flax_key_tuple, flax_tensor in flax_state_dict.items():
      UpperCAmelCase							=     flax_key_tuple.split("""."""   )

      if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
         UpperCAmelCase							=     flax_key_tuple_array[:-1] + ["""weight"""]
         UpperCAmelCase							=     jnp.transpose(lowerCAmelCase_  ,						(3, 2, 0, 1)   )
      elif flax_key_tuple_array[-1] == "kernel":
         UpperCAmelCase							=     flax_key_tuple_array[:-1] + ["""weight"""]
         UpperCAmelCase							=     flax_tensor.T
      elif flax_key_tuple_array[-1] == "scale":
         UpperCAmelCase							=     flax_key_tuple_array[:-1] + ["""weight"""]

      if "time_embedding" not in flax_key_tuple_array:
         for i, flax_key_tuple_string in enumerate(lowerCAmelCase_   ):
            UpperCAmelCase							=     (
                flax_key_tuple_string.replace("""_0"""  ,						""".0"""   )
                .replace("""_1"""  ,						""".1"""   )
                .replace("""_2"""  ,						""".2"""   )
                .replace("""_3"""  ,						""".3"""   )
                .replace("""_4"""  ,						""".4"""   )
                .replace("""_5"""  ,						""".5"""   )
                .replace("""_6"""  ,						""".6"""   )
                .replace("""_7"""  ,						""".7"""   )
                .replace("""_8"""  ,						""".8"""   )
                .replace("""_9"""  ,						""".9"""   )
            )

      UpperCAmelCase							=     """.""".join(lowerCAmelCase_   )

      if flax_key in pt_model_dict:
         if flax_tensor.shape != pt_model_dict[flax_key].shape:
            raise ValueError(
                F"""Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected """
                F"""to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."""   )
         else:
            # add weight to pytorch dict
            UpperCAmelCase							=     np.asarray(lowerCAmelCase_   ) if not isinstance(lowerCAmelCase_  ,						np.ndarray   ) else flax_tensor
            UpperCAmelCase							=     torch.from_numpy(lowerCAmelCase_   )
            # remove from missing keys
            missing_keys.remove(lowerCAmelCase_   )
      else:
         # weight is not expected by PyTorch model
         unexpected_keys.append(lowerCAmelCase_   )

   pt_model.load_state_dict(lowerCAmelCase_   )

   # re-transform missing_keys to list
   UpperCAmelCase							=     list(lowerCAmelCase_   )

   if len(lowerCAmelCase_   ) > 0:
      logger.warning(
          """Some weights of the Flax model were not used when initializing the PyTorch model"""
          F""" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"""
          F""" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"""
          """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
          F""" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"""
          """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
          """ FlaxBertForSequenceClassification model)."""   )
   if len(lowerCAmelCase_   ) > 0:
      logger.warning(
          F"""Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"""
          F""" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"""
          """ use it for predictions and inference."""   )

   return pt_model



