





'''simple docstring'''


import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

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


if is_torch_available():
		import torch

		from transformers import (
		    MODEL_FOR_PRETRAINING_MAPPING,
		    NezhaForMaskedLM,
		    NezhaForMultipleChoice,
		    NezhaForNextSentencePrediction,
		    NezhaForPreTraining,
		    NezhaForQuestionAnswering,
		    NezhaForSequenceClassification,
		    NezhaForTokenClassification,
		    NezhaModel,
		)
		from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST






class       UpperCamelCase__  :







	def __init__(				self			:   List[Any] , lowerCamelCase			:   Optional[Any] , lowerCamelCase			:   List[str]=1_3 , lowerCamelCase			:   Optional[int]=7 , lowerCamelCase			:   str=True , lowerCamelCase			:   Union[str, Any]=True , lowerCamelCase			:   Dict=True , lowerCamelCase			:   Optional[int]=True , lowerCamelCase			:   Any=9_9 , lowerCamelCase			:   str=3_2 , lowerCamelCase			:   Dict=5 , lowerCamelCase			:   Dict=4 , lowerCamelCase			:   Optional[int]=3_7 , lowerCamelCase			:   Tuple="gelu" , lowerCamelCase			:   Optional[Any]=0.1 , lowerCamelCase			:   Any=0.1 , lowerCamelCase			:   Tuple=1_2_8 , lowerCamelCase			:   Union[str, Any]=3_2 , lowerCamelCase			:   str=1_6 , lowerCamelCase			:   int=2 , lowerCamelCase			:   Optional[int]=0.02 , lowerCamelCase			:   List[str]=3 , lowerCamelCase			:   List[str]=4 , lowerCamelCase			:   Optional[Any]=None , ):




				'''simple docstring'''
				a__     = parent
				a__     = batch_size
				a__     = seq_length
				a__     = is_training
				a__     = use_input_mask
				a__     = use_token_type_ids
				a__     = use_labels
				a__     = vocab_size
				a__     = hidden_size
				a__     = num_hidden_layers
				a__     = num_attention_heads
				a__     = intermediate_size
				a__     = hidden_act
				a__     = hidden_dropout_prob
				a__     = attention_probs_dropout_prob
				a__     = max_position_embeddings
				a__     = type_vocab_size
				a__     = type_sequence_label_size
				a__     = initializer_range
				a__     = num_labels
				a__     = num_choices
				a__     = scope







	def 						__a   (				self			:   Optional[Any]		):




				'''simple docstring'''
				a__     = ids_tensor([self.batch_size, self.seq_length] , self.vocab_size		)

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

				a__     = None
				if self.use_token_type_ids:
							a__     = ids_tensor([self.batch_size, self.seq_length] , self.type_vocab_size		)

				a__     = None
				a__     = None
				a__     = None
				if self.use_labels:
							a__     = ids_tensor([self.batch_size] , self.type_sequence_label_size		)
							a__     = ids_tensor([self.batch_size, self.seq_length] , self.num_labels		)
							a__     = ids_tensor([self.batch_size] , self.num_choices		)

				a__     = self.get_config()

				return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







	def 						__a   (				self			:   Tuple		):




				'''simple docstring'''
				return NezhaConfig(
				    vocab_size=self.vocab_size , hidden_size=self.hidden_size , num_hidden_layers=self.num_hidden_layers , num_attention_heads=self.num_attention_heads , intermediate_size=self.intermediate_size , hidden_act=self.hidden_act , hidden_dropout_prob=self.hidden_dropout_prob , attention_probs_dropout_prob=self.attention_probs_dropout_prob , max_position_embeddings=self.max_position_embeddings , type_vocab_size=self.type_vocab_size , is_decoder=lowerCamelCase , initializer_range=self.initializer_range , )







	def 						__a   (				self			:   Optional[int]		):




				'''simple docstring'''
				(
				    (
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				
				)     = self.prepare_config_and_inputs()

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

				return (
				    config,
				    input_ids,
				    token_type_ids,
				    input_mask,
				    sequence_labels,
				    token_labels,
				    choice_labels,
				    encoder_hidden_states,
				    encoder_attention_mask,
				)







	def 						__a   (				self			:   Optional[Any] , lowerCamelCase			:   Optional[Any] , lowerCamelCase			:   List[Any] , lowerCamelCase			:   str , lowerCamelCase			:   str , lowerCamelCase			:   int , lowerCamelCase			:   str , lowerCamelCase			:   Tuple		):




				'''simple docstring'''
				a__     = NezhaModel(config=lowerCamelCase		)
				model.to(lowerCamelCase		)
				model.eval()
				a__     = model(lowerCamelCase , attention_mask=lowerCamelCase , token_type_ids=lowerCamelCase		)
				a__     = model(lowerCamelCase , token_type_ids=lowerCamelCase		)
				a__     = model(lowerCamelCase		)
				self.parent.assertEqual(result.last_hidden_state.shape , (self.batch_size, self.seq_length, self.hidden_size)		)
				self.parent.assertEqual(result.pooler_output.shape , (self.batch_size, self.hidden_size)		)







	def 						__a   (				self			:   str , lowerCamelCase			:   Tuple , lowerCamelCase			:   Dict , lowerCamelCase			:   str , lowerCamelCase			:   Dict , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   str , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   Dict , lowerCamelCase			:   int , ):




				'''simple docstring'''
				a__     = True
				a__     = NezhaModel(lowerCamelCase		)
				model.to(lowerCamelCase		)
				model.eval()
				a__     = model(
				    lowerCamelCase , attention_mask=lowerCamelCase , token_type_ids=lowerCamelCase , encoder_hidden_states=lowerCamelCase , encoder_attention_mask=lowerCamelCase , )
				a__     = model(
				    lowerCamelCase , attention_mask=lowerCamelCase , token_type_ids=lowerCamelCase , encoder_hidden_states=lowerCamelCase , )
				a__     = model(lowerCamelCase , attention_mask=lowerCamelCase , token_type_ids=lowerCamelCase		)
				self.parent.assertEqual(result.last_hidden_state.shape , (self.batch_size, self.seq_length, self.hidden_size)		)
				self.parent.assertEqual(result.pooler_output.shape , (self.batch_size, self.hidden_size)		)







	def 						__a   (				self			:   Optional[Any] , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   Tuple , lowerCamelCase			:   Any , lowerCamelCase			:   List[Any] , lowerCamelCase			:   List[str] , lowerCamelCase			:   Tuple , lowerCamelCase			:   int		):




				'''simple docstring'''
				a__     = NezhaForMaskedLM(config=lowerCamelCase		)
				model.to(lowerCamelCase		)
				model.eval()
				a__     = model(lowerCamelCase , attention_mask=lowerCamelCase , token_type_ids=lowerCamelCase , labels=lowerCamelCase		)
				self.parent.assertEqual(result.logits.shape , (self.batch_size, self.seq_length, self.vocab_size)		)







	def 						__a   (				self			:   Optional[int] , lowerCamelCase			:   str , lowerCamelCase			:   Dict , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   List[str] , lowerCamelCase			:   str , lowerCamelCase			:   List[Any] , lowerCamelCase			:   Union[str, Any]		):




				'''simple docstring'''
				a__     = NezhaForNextSentencePrediction(config=lowerCamelCase		)
				model.to(lowerCamelCase		)
				model.eval()
				a__     = model(
				    lowerCamelCase , attention_mask=lowerCamelCase , token_type_ids=lowerCamelCase , labels=lowerCamelCase , )
				self.parent.assertEqual(result.logits.shape , (self.batch_size, 2)		)







	def 						__a   (				self			:   Optional[int] , lowerCamelCase			:   Optional[Any] , lowerCamelCase			:   Optional[int] , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   str , lowerCamelCase			:   List[str] , lowerCamelCase			:   Optional[Any] , lowerCamelCase			:   Tuple		):




				'''simple docstring'''
				a__     = NezhaForPreTraining(config=lowerCamelCase		)
				model.to(lowerCamelCase		)
				model.eval()
				a__     = model(
				    lowerCamelCase , attention_mask=lowerCamelCase , token_type_ids=lowerCamelCase , labels=lowerCamelCase , next_sentence_label=lowerCamelCase , )
				self.parent.assertEqual(result.prediction_logits.shape , (self.batch_size, self.seq_length, self.vocab_size)		)
				self.parent.assertEqual(result.seq_relationship_logits.shape , (self.batch_size, 2)		)







	def 						__a   (				self			:   List[str] , lowerCamelCase			:   int , lowerCamelCase			:   str , lowerCamelCase			:   Optional[int] , lowerCamelCase			:   str , lowerCamelCase			:   Tuple , lowerCamelCase			:   List[str] , lowerCamelCase			:   Tuple		):




				'''simple docstring'''
				a__     = NezhaForQuestionAnswering(config=lowerCamelCase		)
				model.to(lowerCamelCase		)
				model.eval()
				a__     = model(
				    lowerCamelCase , attention_mask=lowerCamelCase , token_type_ids=lowerCamelCase , start_positions=lowerCamelCase , end_positions=lowerCamelCase , )
				self.parent.assertEqual(result.start_logits.shape , (self.batch_size, self.seq_length)		)
				self.parent.assertEqual(result.end_logits.shape , (self.batch_size, self.seq_length)		)







	def 						__a   (				self			:   List[str] , lowerCamelCase			:   int , lowerCamelCase			:   Tuple , lowerCamelCase			:   List[str] , lowerCamelCase			:   List[Any] , lowerCamelCase			:   str , lowerCamelCase			:   List[Any] , lowerCamelCase			:   Tuple		):




				'''simple docstring'''
				a__     = self.num_labels
				a__     = NezhaForSequenceClassification(lowerCamelCase		)
				model.to(lowerCamelCase		)
				model.eval()
				a__     = model(lowerCamelCase , attention_mask=lowerCamelCase , token_type_ids=lowerCamelCase , labels=lowerCamelCase		)
				self.parent.assertEqual(result.logits.shape , (self.batch_size, self.num_labels)		)







	def 						__a   (				self			:   Optional[int] , lowerCamelCase			:   List[str] , lowerCamelCase			:   List[Any] , lowerCamelCase			:   Optional[Any] , lowerCamelCase			:   Tuple , lowerCamelCase			:   str , lowerCamelCase			:   List[str] , lowerCamelCase			:   int		):




				'''simple docstring'''
				a__     = self.num_labels
				a__     = NezhaForTokenClassification(config=lowerCamelCase		)
				model.to(lowerCamelCase		)
				model.eval()
				a__     = model(lowerCamelCase , attention_mask=lowerCamelCase , token_type_ids=lowerCamelCase , labels=lowerCamelCase		)
				self.parent.assertEqual(result.logits.shape , (self.batch_size, self.seq_length, self.num_labels)		)







	def 						__a   (				self			:   Dict , lowerCamelCase			:   Tuple , lowerCamelCase			:   List[Any] , lowerCamelCase			:   Any , lowerCamelCase			:   Optional[Any] , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   str , lowerCamelCase			:   Optional[int]		):




				'''simple docstring'''
				a__     = self.num_choices
				a__     = NezhaForMultipleChoice(config=lowerCamelCase		)
				model.to(lowerCamelCase		)
				model.eval()
				a__     = input_ids.unsqueeze(1		).expand(-1 , self.num_choices , -1		).contiguous()
				a__     = token_type_ids.unsqueeze(1		).expand(-1 , self.num_choices , -1		).contiguous()
				a__     = input_mask.unsqueeze(1		).expand(-1 , self.num_choices , -1		).contiguous()
				a__     = model(
				    lowerCamelCase , attention_mask=lowerCamelCase , token_type_ids=lowerCamelCase , labels=lowerCamelCase , )
				self.parent.assertEqual(result.logits.shape , (self.batch_size, self.num_choices)		)







	def 						__a   (				self			:   Any		):




				'''simple docstring'''
				a__     = self.prepare_config_and_inputs()
				(
				    (
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				
				)     = config_and_inputs
				a__     = {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
				return config, inputs_dict






@require_torch
class       UpperCamelCase__  (		__lowerCAmelCase	,__lowerCAmelCase	,__lowerCAmelCase	,unittest.TestCase     ):
	lowerCAmelCase__					:			Any								=						(
	    (
	        NezhaModel,
	        NezhaForMaskedLM,
	        NezhaForMultipleChoice,
	        NezhaForNextSentencePrediction,
	        NezhaForPreTraining,
	        NezhaForQuestionAnswering,
	        NezhaForSequenceClassification,
	        NezhaForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	lowerCAmelCase__					:			Union[str, Any]								=						(
	    {
	        "feature-extraction": NezhaModel,
	        "fill-mask": NezhaForMaskedLM,
	        "question-answering": NezhaForQuestionAnswering,
	        "text-classification": NezhaForSequenceClassification,
	        "token-classification": NezhaForTokenClassification,
	        "zero-shot": NezhaForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	lowerCAmelCase__					:			str								=						True







	def 						__a   (				self			:   Union[str, Any] , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   Optional[Any] , lowerCamelCase			:   int=False		):




				'''simple docstring'''
				a__     = super()._prepare_for_class(lowerCamelCase , lowerCamelCase , return_labels=lowerCamelCase		)

				if return_labels:
							if model_class in get_values(lowerCamelCase		):
										a__     = torch.zeros(
										    (self.model_tester.batch_size, self.model_tester.seq_length) , dtype=torch.long , device=lowerCamelCase		)
										a__     = torch.zeros(
										    self.model_tester.batch_size , dtype=torch.long , device=lowerCamelCase		)
				return inputs_dict







	def 						__a   (				self			:   str		):




				'''simple docstring'''
				a__     = NezhaModelTester(self		)
				a__     = ConfigTester(self , config_class=lowerCamelCase , hidden_size=3_7		)







	def 						__a   (				self			:   Optional[int]		):




				'''simple docstring'''
				self.config_tester.run_common_tests()







	def 						__a   (				self			:   List[Any]		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*lowerCamelCase		)







	def 						__a   (				self			:   str		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs_for_decoder()
				self.model_tester.create_and_check_model_as_decoder(*lowerCamelCase		)







	def 						__a   (				self			:   Union[str, Any]		):




				'''simple docstring'''
				# This regression test was failing with PyTorch < 1.3
				(
				    (
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				(
				    a__
				)   ,				
				)     = self.model_tester.prepare_config_and_inputs_for_decoder()

				a__     = None

				self.model_tester.create_and_check_model_as_decoder(
				    lowerCamelCase , lowerCamelCase , lowerCamelCase , lowerCamelCase , lowerCamelCase , lowerCamelCase , lowerCamelCase , lowerCamelCase , lowerCamelCase , )







	def 						__a   (				self			:   List[str]		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase		)







	def 						__a   (				self			:   Dict		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_multiple_choice(*lowerCamelCase		)







	def 						__a   (				self			:   Any		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_next_sequence_prediction(*lowerCamelCase		)







	def 						__a   (				self			:   List[Any]		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_pretraining(*lowerCamelCase		)







	def 						__a   (				self			:   List[str]		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_question_answering(*lowerCamelCase		)







	def 						__a   (				self			:   List[str]		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_sequence_classification(*lowerCamelCase		)







	def 						__a   (				self			:   Union[str, Any]		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*lowerCamelCase		)







	@slow
	def 						__a   (				self			:   List[Any]		):




				'''simple docstring'''
				for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							a__     = NezhaModel.from_pretrained(lowerCamelCase		)
							self.assertIsNotNone(lowerCamelCase		)







	@slow
	@require_torch_gpu
	def 						__a   (				self			:   Any		):




				'''simple docstring'''
				a__   ,				a__     = self.model_tester.prepare_config_and_inputs_for_common()
				for model_class in self.all_model_classes:
							# NezhaForMultipleChoice behaves incorrectly in JIT environments.
							if model_class == NezhaForMultipleChoice:
										return

							a__     = True
							a__     = model_class(config=lowerCamelCase		)

							a__     = self._prepare_for_class(lowerCamelCase , lowerCamelCase		)
							a__     = torch.jit.trace(
							    lowerCamelCase , (inputs_dict["input_ids"].to("cpu"		), inputs_dict["attention_mask"].to("cpu"		))		)

							with tempfile.TemporaryDirectory() as tmp:
										torch.jit.save(lowerCamelCase , os.path.join(lowerCamelCase , "bert.pt"		)		)
										a__     = torch.jit.load(os.path.join(lowerCamelCase , "bert.pt"		) , map_location=lowerCamelCase		)
										loaded(inputs_dict["input_ids"].to(lowerCamelCase		) , inputs_dict["attention_mask"].to(lowerCamelCase		)		)






@require_torch
class       UpperCamelCase__  (		unittest.TestCase     ):







	@slow
	def 						__a   (				self			:   Dict		):




				'''simple docstring'''
				a__     = NezhaModel.from_pretrained("sijunhe/nezha-cn-base"		)
				a__     = torch.tensor([[0, 1, 2, 3, 4, 5]]		)
				a__     = torch.tensor([[0, 1, 1, 1, 1, 1]]		)
				with torch.no_grad():
							a__     = model(lowerCamelCase , attention_mask=lowerCamelCase		)[0]
				a__     = torch.Size((1, 6, 7_6_8)		)
				self.assertEqual(output.shape , lowerCamelCase		)
				a__     = torch.tensor([[[0.0685, 0.2441, 0.1102], [0.0600, 0.1906, 0.1349], [0.0221, 0.0819, 0.0586]]]		)

				self.assertTrue(torch.allclose(output[:, 1:4, 1:4] , lowerCamelCase , atol=1e-4		)		)







	@slow
	def 						__a   (				self			:   Tuple		):




				'''simple docstring'''
				a__     = NezhaForMaskedLM.from_pretrained("sijunhe/nezha-cn-base"		)
				a__     = torch.tensor([[0, 1, 2, 3, 4, 5]]		)
				a__     = torch.tensor([[1, 1, 1, 1, 1, 1]]		)
				with torch.no_grad():
							a__     = model(lowerCamelCase , attention_mask=lowerCamelCase		)[0]
				a__     = torch.Size((1, 6, 2_1_1_2_8)		)
				self.assertEqual(output.shape , lowerCamelCase		)
				a__     = torch.tensor(
				    [[-2.7939, -1.7902, -2.2189], [-2.8585, -1.8908, -2.3723], [-2.6499, -1.7750, -2.2558]]		)

				self.assertTrue(torch.allclose(output[:, 1:4, 1:4] , lowerCamelCase , atol=1e-4		)		)










'''simple docstring'''


import os
from pickle import UnpicklingError
from typing import Dict, Tuple

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

import transformers

from .utils import logging


lowerCAmelCase_   :   Union[str, Any]       =					logging.get_logger(__name__)







def    _lowerCamelCase  (__lowerCamelCase  :     List[Any]					,  __lowerCamelCase  :     Optional[int]					,  __lowerCamelCase  :     List[str]					,  __lowerCamelCase  :     Union[str, Any]=False     )   ->    Dict:
			try:
						import torch  # noqa: F401
			except ImportError:
						logger.error(
						    "Loading a PyTorch model in Flax, 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

			if not is_sharded:
						a__     = os.path.abspath(__lowerCamelCase     )
						logger.info(f'''Loading PyTorch weights from {pt_path}'''     )

						a__     = torch.load(__lowerCamelCase					,  map_location="cpu"     )
						logger.info(f'''PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()     ):,} parameters.'''     )

						a__     = convert_pytorch_state_dict_to_flax(__lowerCamelCase					,  __lowerCamelCase     )
			else:
						# model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
						a__     = convert_pytorch_sharded_state_dict_to_flax(__lowerCamelCase					,  __lowerCamelCase     )
			return flax_state_dict







def    _lowerCamelCase  (__lowerCamelCase  :     Tuple[str]					,  __lowerCamelCase  :     np.ndarray					,  __lowerCamelCase  :     Dict[str, jnp.ndarray]					,  __lowerCamelCase  :     str					,  )   ->    (Tuple[str], np.ndarray):

			def is_key_or_prefix_key_in_dict(__lowerCamelCase  :     Tuple[str]     ) -> bool:
						return len(set(__lowerCamelCase     ) & {key, (model_prefix,) + key}     ) > 0

			# layer norm
			a__     = pt_tuple_key[:-1] + ("scale",)
			if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						return renamed_pt_tuple_key, pt_tensor

			# batch norm layer mean
			a__     = pt_tuple_key[:-1] + ("mean",)
			if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						return renamed_pt_tuple_key, pt_tensor

			# batch norm layer var
			a__     = pt_tuple_key[:-1] + ("var",)
			if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						return renamed_pt_tuple_key, pt_tensor

			# embedding
			a__     = pt_tuple_key[:-1] + ("embedding",)
			if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						return renamed_pt_tuple_key, pt_tensor

			# conv layer
			a__     = pt_tuple_key[:-1] + ("kernel",)
			if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						a__     = pt_tensor.transpose(2					,  3					,  1					,  0     )
						return renamed_pt_tuple_key, pt_tensor

			# linear layer
			a__     = pt_tuple_key[:-1] + ("kernel",)
			if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						a__     = pt_tensor.T
						return renamed_pt_tuple_key, pt_tensor

			# old PyTorch layer norm weight
			a__     = pt_tuple_key[:-1] + ("weight",)
			if pt_tuple_key[-1] == "gamma":
						return renamed_pt_tuple_key, pt_tensor

			# old PyTorch layer norm bias
			a__     = pt_tuple_key[:-1] + ("bias",)
			if pt_tuple_key[-1] == "beta":
						return renamed_pt_tuple_key, pt_tensor

			# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
			a__     = None
			if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
						a__     = pt_tuple_key[-2] + "_g"
			elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
						a__     = pt_tuple_key[-2] + "_v"
			if name is not None:
						a__     = pt_tuple_key[:-3] + (name,)
						return renamed_pt_tuple_key, pt_tensor

			return pt_tuple_key, pt_tensor







def    _lowerCamelCase  (__lowerCamelCase  :     int					,  __lowerCamelCase  :     List[Any]     )   ->    Any:
			# convert pytorch tensor to numpy
			a__     = {k: v.numpy() for k, v in pt_state_dict.items()}

			a__     = flax_model.base_model_prefix

			# use params dict if the model contains batch norm layers
			if "params" in flax_model.params:
						a__     = flax_model.params["params"]
			else:
						a__     = flax_model.params
			a__     = flatten_dict(__lowerCamelCase     )

			# add batch_stats keys,values to dict
			if "batch_stats" in flax_model.params:
						a__     = flatten_dict(flax_model.params["batch_stats"]     )
						random_flax_state_dict.update(__lowerCamelCase     )

			a__     = {}

			a__     = (model_prefix not in flax_model_params) and (
			    model_prefix in {k.split("."     )[0] for k in pt_state_dict.keys()}
			)
			a__     = (model_prefix in flax_model_params) and (
			    model_prefix not in {k.split("."     )[0] for k in pt_state_dict.keys()}
			)

			# Need to change some parameters name to match Flax names
			for pt_key, pt_tensor in pt_state_dict.items():
						a__     = tuple(pt_key.split("."     )     )

						# remove base model prefix if necessary
						a__     = pt_tuple_key[0] == model_prefix
						if load_model_with_head_into_base_model and has_base_model_prefix:
									a__     = pt_tuple_key[1:]

						# Correctly rename weight parameters
						a__   ,				a__     = rename_key_and_reshape_tensor(
						    __lowerCamelCase					,  __lowerCamelCase					,  __lowerCamelCase					,  __lowerCamelCase     )

						# add model prefix if necessary
						a__     = (model_prefix,) + flax_key in random_flax_state_dict
						if load_base_model_into_model_with_head and require_base_model_prefix:
									a__     = (model_prefix,) + flax_key

						if flax_key in random_flax_state_dict:
									if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
												raise ValueError(
												    f'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
												    f'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''     )

        # add batch stats if the model contains batchnorm layers
						if "batch_stats" in flax_model.params:
									if "mean" in flax_key[-1] or "var" in flax_key[-1]:
												a__     = jnp.asarray(__lowerCamelCase     )
												continue
									# remove num_batches_tracked key
									if "num_batches_tracked" in flax_key[-1]:
												flax_state_dict.pop(__lowerCamelCase					,  __lowerCamelCase     )
												continue

									# also add unexpected weight so that warning is thrown
									a__     = jnp.asarray(__lowerCamelCase     )

						else:
									# also add unexpected weight so that warning is thrown
									a__     = jnp.asarray(__lowerCamelCase     )

			return unflatten_dict(__lowerCamelCase     )







def    _lowerCamelCase  (__lowerCamelCase  :     Union[str, Any]					,  __lowerCamelCase  :     Any     )   ->    Union[str, Any]:
			import torch

			# Load the index
			a__     = {}
			for shard_file in shard_filenames:
						# load using msgpack utils
						a__     = torch.load(__lowerCamelCase     )
						a__     = {k: v.numpy() for k, v in pt_state_dict.items()}

						a__     = flax_model.base_model_prefix

						# use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
						if "batch_stats" in flax_model.params:
									a__     = flax_model.params["params"]

									a__     = flatten_dict(__lowerCamelCase     )
									random_flax_state_dict.update(flatten_dict(flax_model.params["batch_stats"]     )     )
						else:
									a__     = flax_model.params
									a__     = flatten_dict(__lowerCamelCase     )

						a__     = (model_prefix not in flax_model_params) and (
						    model_prefix in {k.split("."     )[0] for k in pt_state_dict.keys()}
						)
						a__     = (model_prefix in flax_model_params) and (
						    model_prefix not in {k.split("."     )[0] for k in pt_state_dict.keys()}
						)
						# Need to change some parameters name to match Flax names
						for pt_key, pt_tensor in pt_state_dict.items():
									a__     = tuple(pt_key.split("."     )     )

									# remove base model prefix if necessary
									a__     = pt_tuple_key[0] == model_prefix
									if load_model_with_head_into_base_model and has_base_model_prefix:
												a__     = pt_tuple_key[1:]

									# Correctly rename weight parameters
									a__   ,				a__     = rename_key_and_reshape_tensor(
									    __lowerCamelCase					,  __lowerCamelCase					,  __lowerCamelCase					,  __lowerCamelCase     )
									# add model prefix if necessary
									a__     = (model_prefix,) + flax_key in random_flax_state_dict
									if load_base_model_into_model_with_head and require_base_model_prefix:
												a__     = (model_prefix,) + flax_key

									if flax_key in random_flax_state_dict:
												if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
															raise ValueError(
															    f'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
															    f'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''     )

            # add batch stats if the model contains batchnorm layers
									if "batch_stats" in flax_model.params:
												if "mean" in flax_key[-1]:
															a__     = jnp.asarray(__lowerCamelCase     )
															continue
												if "var" in flax_key[-1]:
															a__     = jnp.asarray(__lowerCamelCase     )
															continue
												# remove num_batches_tracked key
												if "num_batches_tracked" in flax_key[-1]:
															flax_state_dict.pop(__lowerCamelCase					,  __lowerCamelCase     )
															continue

												# also add unexpected weight so that warning is thrown
												a__     = jnp.asarray(__lowerCamelCase     )

									else:
												# also add unexpected weight so that warning is thrown
												a__     = jnp.asarray(__lowerCamelCase     )
			return unflatten_dict(__lowerCamelCase     )







def    _lowerCamelCase  (__lowerCamelCase  :     Tuple					,  __lowerCamelCase  :     Tuple     )   ->    Any:
			a__     = os.path.abspath(__lowerCamelCase     )
			logger.info(f'''Loading Flax weights from {flax_checkpoint_path}'''     )

			# import correct flax class
			a__     = getattr(__lowerCamelCase					,  "Flax" + model.__class__.__name__     )

			# load flax weight dict
			with open(__lowerCamelCase					,  "rb"     ) as state_f:
						try:
									a__     = from_bytes(__lowerCamelCase					,  state_f.read()     )
						except UnpicklingError:
									raise EnvironmentError(f'''Unable to convert {flax_checkpoint_path} to Flax deserializable object. '''     )

			return load_flax_weights_in_pytorch_model(__lowerCamelCase					,  __lowerCamelCase     )







def    _lowerCamelCase  (__lowerCamelCase  :     str					,  __lowerCamelCase  :     Tuple     )   ->    Dict:

			try:
						import torch  # noqa: F401
			except ImportError:
						logger.error(
						    "Loading a 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
			a__     = flatten_dict(jax.tree_util.tree_map(lambda __lowerCamelCase     : x.dtype == jnp.bfloataa					,  __lowerCamelCase     )     ).values()
			if any(__lowerCamelCase     ):
						# convert all weights to fp32 if the 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."     )
						a__     = jax.tree_util.tree_map(
						    lambda __lowerCamelCase     : params.astype(np.floataa     ) if params.dtype == jnp.bfloataa else params					,  __lowerCamelCase     )

			a__     = flatten_dict(__lowerCamelCase     )
			a__     = pt_model.state_dict()

			a__     = (pt_model.base_model_prefix in flax_state) and (
			    pt_model.base_model_prefix not in {k.split("."     )[0] for k in pt_model_dict.keys()}
			)
			a__     = (pt_model.base_model_prefix not in flax_state) and (
			    pt_model.base_model_prefix in {k.split("."     )[0] for k in pt_model_dict.keys()}
			)

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

			for flax_key_tuple, flax_tensor in flax_state_dict.items():
						a__     = flax_key_tuple[0] == pt_model.base_model_prefix
						a__     = ".".join((pt_model.base_model_prefix,) + flax_key_tuple     ) in pt_model_dict

						# adapt flax_key to prepare for loading from/to base model only
						if load_model_with_head_into_base_model and has_base_model_prefix:
									a__     = flax_key_tuple[1:]
						elif load_base_model_into_model_with_head and require_base_model_prefix:
									a__     = (pt_model.base_model_prefix,) + flax_key_tuple

						# rename flax weights to PyTorch format
						if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(__lowerCamelCase     ) not in pt_model_dict:
									# conv layer
									a__     = flax_key_tuple[:-1] + ("weight",)
									a__     = jnp.transpose(__lowerCamelCase					,  (3, 2, 0, 1)     )
						elif flax_key_tuple[-1] == "kernel" and ".".join(__lowerCamelCase     ) not in pt_model_dict:
									# linear layer
									a__     = flax_key_tuple[:-1] + ("weight",)
									a__     = flax_tensor.T
						elif flax_key_tuple[-1] in ["scale", "embedding"]:
									a__     = flax_key_tuple[:-1] + ("weight",)

						# adding batch stats from flax batch norm to pt
						elif "mean" in flax_key_tuple[-1]:
									a__     = flax_key_tuple[:-1] + ("running_mean",)
						elif "var" in flax_key_tuple[-1]:
									a__     = flax_key_tuple[:-1] + ("running_var",)

						if "batch_stats" in flax_state:
									a__     = ".".join(flax_key_tuple[1:]     )  # Remove the params/batch_stats header
						else:
									a__     = ".".join(__lowerCamelCase     )

						# We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
						a__     = {}
						# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
						for key in pt_model_dict:
									a__     = key.split("."     )
									a__     = None
									if key_components[-3::2] == ["parametrizations", "original0"]:
												a__     = key_components[-2] + "_g"
									elif key_components[-3::2] == ["parametrizations", "original1"]:
												a__     = key_components[-2] + "_v"
									if name is not None:
												a__     = key_components[:-3] + [name]
												a__     = ".".join(__lowerCamelCase     )
												a__     = key

						if flax_key in special_pt_names:
									a__     = special_pt_names[flax_key]

						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
												a__     = np.asarray(__lowerCamelCase     ) if not isinstance(__lowerCamelCase					,  np.ndarray     ) else flax_tensor
												a__     = 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
			a__     = 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)."     )
			else:
						logger.warning(f'''All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n'''     )
			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."     )
			else:
						logger.warning(
						    f'''All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n'''
						    "If your task is similar to the task the model of the checkpoint was trained on, "
						    f'''you can already use {pt_model.__class__.__name__} for predictions without further training.'''     )

			return pt_model



