





"""simple docstring"""






from __future__ import annotations

import unittest

from transformers import FunnelConfig, is_tf_available
from transformers.testing_utils import require_tf

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import tensorflow as tf

		from transformers import (
		    TFFunnelBaseModel,
		    TFFunnelForMaskedLM,
		    TFFunnelForMultipleChoice,
		    TFFunnelForPreTraining,
		    TFFunnelForQuestionAnswering,
		    TFFunnelForSequenceClassification,
		    TFFunnelForTokenClassification,
		    TFFunnelModel,
		)



class 					__UpperCAmelCase							:
				def __init__(		self    :			Dict	, a_    :			int	, a_    :			Any=13	, a_    :			List[Any]=7	, a_    :			Dict=True	, a_    :			Optional[Any]=True	, a_    :			Optional[int]=True	, a_    :			List[str]=True	, a_    :			int=99	, a_    :			List[str]=[1, 1, 2]	, a_    :			Any=1	, a_    :			List[Any]=32	, a_    :			str=4	, a_    :			Union[str, Any]=8	, a_    :			List[str]=37	, a_    :			str="gelu_new"	, a_    :			Optional[int]=0.1	, a_    :			Union[str, Any]=0.1	, a_    :			Tuple=0.0	, a_    :			str=5_12	, a_    :			Optional[int]=3	, a_    :			List[str]=0.02	, a_    :			List[Any]=3	, a_    :			Dict=4	, a_    :			List[str]=None	, a_    :			List[str]=False	, )	->	List[Any]:



							'''simple docstring'''



							a__  :	Any		 =			parent
							a__  :	Optional[int]		 =			batch_size
							a__  :	List[Any]		 =			seq_length
							a__  :	Optional[Any]		 =			is_training
							a__  :	Optional[Any]		 =			use_input_mask
							a__  :	int		 =			use_token_type_ids
							a__  :	Union[str, Any]		 =			use_labels
							a__  :	Optional[int]		 =			vocab_size
							a__  :	Dict		 =			block_sizes
							a__  :	Any		 =			num_decoder_layers
							a__  :	Any		 =			d_model
							a__  :	List[Any]		 =			n_head
							a__  :	str		 =			d_head
							a__  :	Union[str, Any]		 =			d_inner
							a__  :	str		 =			hidden_act
							a__  :	Any		 =			hidden_dropout
							a__  :	Union[str, Any]		 =			attention_dropout
							a__  :	Optional[Any]		 =			activation_dropout
							a__  :	Any		 =			max_position_embeddings
							a__  :	Any		 =			type_vocab_size
							a__  :	List[str]		 =			2
							a__  :	Optional[Any]		 =			num_labels
							a__  :	Dict		 =			num_choices
							a__  :	Union[str, Any]		 =			scope
							a__  :	str		 =			initializer_std

							# Used in the tests to check the size of the first attention layer
							a__  :	Optional[int]		 =			n_head
							# Used in the tests to check the size of the first hidden state
							a__  :	Dict		 =			self.d_model
							# Used in the tests to check the number of output hidden states/attentions
							a__  :	Tuple		 =			sum(self.block_sizes    ) + (0 if base else self.num_decoder_layers)
							# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
							# the last hidden state of the first block (which is the first hidden state of the decoder).
							if not base:
										a__  :	int		 =			self.num_hidden_layers + 2
				def       UpperCAmelCase      (		self    :			Dict    )	->	Dict:



							'''simple docstring'''



							a__  :	Dict		 =			ids_tensor([self.batch_size, self.seq_length]	, self.vocab_size    )

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

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

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

							a__  :	Any		 =			FunnelConfig(
							    vocab_size=self.vocab_size	, block_sizes=self.block_sizes	, num_decoder_layers=self.num_decoder_layers	, d_model=self.d_model	, n_head=self.n_head	, d_head=self.d_head	, d_inner=self.d_inner	, hidden_act=self.hidden_act	, hidden_dropout=self.hidden_dropout	, attention_dropout=self.attention_dropout	, activation_dropout=self.activation_dropout	, max_position_embeddings=self.max_position_embeddings	, type_vocab_size=self.type_vocab_size	, initializer_std=self.initializer_std	, )

							return (
							    config,
							    input_ids,
							    token_type_ids,
							    input_mask,
							    sequence_labels,
							    token_labels,
							    choice_labels,
							)
				def       UpperCAmelCase      (		self    :			Dict	, a_    :			Dict	, a_    :			List[Any]	, a_    :			Dict	, a_    :			Dict	, a_    :			Dict	, a_    :			Union[str, Any]	, a_    :			Union[str, Any]	, )	->	List[str]:



							'''simple docstring'''



							a__  :	List[str]		 =			TFFunnelModel(config=a_    )
							a__  :	Dict		 =			{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							a__  :	Dict		 =			model(a_    )

							a__  :	Optional[Any]		 =			[input_ids, input_mask]
							a__  :	Optional[Any]		 =			model(a_    )

							a__  :	str		 =			model(a_    )
							self.parent.assertEqual(result.last_hidden_state.shape	, (self.batch_size, self.seq_length, self.d_model)    )

							a__  :	Tuple		 =			False
							a__  :	List[Any]		 =			TFFunnelModel(config=a_    )
							a__  :	Tuple		 =			model(a_    )
							self.parent.assertEqual(result.last_hidden_state.shape	, (self.batch_size, self.seq_length, self.d_model)    )

							a__  :	List[str]		 =			False
							a__  :	int		 =			TFFunnelModel(config=a_    )
							a__  :	Tuple		 =			model(a_    )
							self.parent.assertEqual(result.last_hidden_state.shape	, (self.batch_size, self.seq_length, self.d_model)    )
				def       UpperCAmelCase      (		self    :			Optional[int]	, a_    :			int	, a_    :			Optional[int]	, a_    :			Optional[Any]	, a_    :			int	, a_    :			Any	, a_    :			int	, a_    :			int	, )	->	Tuple:



							'''simple docstring'''



							a__  :	Optional[int]		 =			TFFunnelBaseModel(config=a_    )
							a__  :	Dict		 =			{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							a__  :	Optional[Any]		 =			model(a_    )

							a__  :	Union[str, Any]		 =			[input_ids, input_mask]
							a__  :	Optional[Any]		 =			model(a_    )

							a__  :	Any		 =			model(a_    )
							self.parent.assertEqual(result.last_hidden_state.shape	, (self.batch_size, 2, self.d_model)    )

							a__  :	Any		 =			False
							a__  :	Dict		 =			TFFunnelBaseModel(config=a_    )
							a__  :	str		 =			model(a_    )
							self.parent.assertEqual(result.last_hidden_state.shape	, (self.batch_size, 3, self.d_model)    )

							a__  :	int		 =			False
							a__  :	str		 =			TFFunnelBaseModel(config=a_    )
							a__  :	Dict		 =			model(a_    )
							self.parent.assertEqual(result.last_hidden_state.shape	, (self.batch_size, 2, self.d_model)    )
				def       UpperCAmelCase      (		self    :			Tuple	, a_    :			str	, a_    :			int	, a_    :			Tuple	, a_    :			List[Any]	, a_    :			Optional[Any]	, a_    :			str	, a_    :			Optional[Any]	, )	->	List[Any]:



							'''simple docstring'''



							a__  :	Optional[int]		 =			TFFunnelForPreTraining(config=a_    )
							a__  :	Dict		 =			{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							a__  :	List[str]		 =			model(a_    )
							self.parent.assertEqual(result.logits.shape	, (self.batch_size, self.seq_length)    )
				def       UpperCAmelCase      (		self    :			str	, a_    :			Optional[Any]	, a_    :			List[str]	, a_    :			List[Any]	, a_    :			Union[str, Any]	, a_    :			Tuple	, a_    :			int	, a_    :			List[str]	, )	->	List[str]:



							'''simple docstring'''



							a__  :	str		 =			TFFunnelForMaskedLM(config=a_    )
							a__  :	Union[str, Any]		 =			{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							a__  :	Tuple		 =			model(a_    )
							self.parent.assertEqual(result.logits.shape	, (self.batch_size, self.seq_length, self.vocab_size)    )
				def       UpperCAmelCase      (		self    :			Any	, a_    :			Tuple	, a_    :			Tuple	, a_    :			Dict	, a_    :			List[Any]	, a_    :			Dict	, a_    :			int	, a_    :			int	, )	->	Tuple:



							'''simple docstring'''



							a__  :	List[Any]		 =			self.num_labels
							a__  :	int		 =			TFFunnelForSequenceClassification(config=a_    )
							a__  :	List[str]		 =			{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							a__  :	List[Any]		 =			model(a_    )
							self.parent.assertEqual(result.logits.shape	, (self.batch_size, self.num_labels)    )
				def       UpperCAmelCase      (		self    :			Any	, a_    :			List[str]	, a_    :			Any	, a_    :			str	, a_    :			str	, a_    :			int	, a_    :			int	, a_    :			List[str]	, )	->	Union[str, Any]:



							'''simple docstring'''



							a__  :	Optional[Any]		 =			self.num_choices
							a__  :	str		 =			TFFunnelForMultipleChoice(config=a_    )
							a__  :	int		 =			tf.tile(tf.expand_dims(a_	, 1    )	, (1, self.num_choices, 1)    )
							a__  :	int		 =			tf.tile(tf.expand_dims(a_	, 1    )	, (1, self.num_choices, 1)    )
							a__  :	Any		 =			tf.tile(tf.expand_dims(a_	, 1    )	, (1, self.num_choices, 1)    )
							a__  :	List[Any]		 =			{
							    "input_ids": multiple_choice_inputs_ids,
							    "attention_mask": multiple_choice_input_mask,
							    "token_type_ids": multiple_choice_token_type_ids,
							}
							a__  :	Any		 =			model(a_    )
							self.parent.assertEqual(result.logits.shape	, (self.batch_size, self.num_choices)    )
				def       UpperCAmelCase      (		self    :			List[str]	, a_    :			Optional[int]	, a_    :			str	, a_    :			str	, a_    :			List[Any]	, a_    :			Any	, a_    :			Optional[Any]	, a_    :			List[str]	, )	->	Tuple:



							'''simple docstring'''



							a__  :	str		 =			self.num_labels
							a__  :	str		 =			TFFunnelForTokenClassification(config=a_    )
							a__  :	List[str]		 =			{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							a__  :	Optional[Any]		 =			model(a_    )
							self.parent.assertEqual(result.logits.shape	, (self.batch_size, self.seq_length, self.num_labels)    )
				def       UpperCAmelCase      (		self    :			Optional[int]	, a_    :			int	, a_    :			int	, a_    :			Optional[Any]	, a_    :			int	, a_    :			Any	, a_    :			str	, a_    :			Tuple	, )	->	List[str]:



							'''simple docstring'''



							a__  :	Optional[int]		 =			TFFunnelForQuestionAnswering(config=a_    )
							a__  :	Any		 =			{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							a__  :	List[str]		 =			model(a_    )
							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       UpperCAmelCase      (		self    :			Any    )	->	Optional[Any]:



							'''simple docstring'''



							a__  :	Any		 =			self.prepare_config_and_inputs()
							(
							    (
							    a__
							) ,	(
							    a__
							) ,	(
							    a__
							) ,	(
							    a__
							) ,	(
							    a__
							) ,	(
							    a__
							) ,	(
							    a__
							) ,	
							)  :	List[Any]		 =			config_and_inputs
							a__  :	Union[str, Any]		 =			{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
							return config, inputs_dict



@require_tf
class 					__UpperCAmelCase							(						_UpperCamelCase						,  _UpperCamelCase						,  unittest.TestCase  ):
				__lowerCamelCase				:	Optional[int]				     =				(
				    (
				        TFFunnelModel,
				        TFFunnelForMaskedLM,
				        TFFunnelForPreTraining,
				        TFFunnelForQuestionAnswering,
				        TFFunnelForTokenClassification,
				    )
				    if is_tf_available()
				    else ()
				)
				__lowerCamelCase				:	Union[str, Any]				     =				(
				    {
				        "feature-extraction": (TFFunnelBaseModel, TFFunnelModel),
				        "fill-mask": TFFunnelForMaskedLM,
				        "question-answering": TFFunnelForQuestionAnswering,
				        "text-classification": TFFunnelForSequenceClassification,
				        "token-classification": TFFunnelForTokenClassification,
				        "zero-shot": TFFunnelForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				__lowerCamelCase				:	List[Any]				     =				False
				__lowerCamelCase				:	Optional[int]				     =				False
				def       UpperCAmelCase      (		self    :			int    )	->	Any:



							'''simple docstring'''



							a__  :	List[Any]		 =			TFFunnelModelTester(self    )
							a__  :	List[str]		 =			ConfigTester(self	, config_class=a_    )
				def       UpperCAmelCase      (		self    :			Tuple    )	->	Optional[int]:



							'''simple docstring'''



							self.config_tester.run_common_tests()
				def       UpperCAmelCase      (		self    :			List[Any]    )	->	Optional[int]:



							'''simple docstring'''



							a__  :	List[str]		 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*a_    )
				def       UpperCAmelCase      (		self    :			Optional[int]    )	->	List[str]:



							'''simple docstring'''



							a__  :	int		 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_pretraining(*a_    )
				def       UpperCAmelCase      (		self    :			Optional[int]    )	->	Optional[Any]:



							'''simple docstring'''



							a__  :	Tuple		 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*a_    )
				def       UpperCAmelCase      (		self    :			Union[str, Any]    )	->	Union[str, Any]:



							'''simple docstring'''



							a__  :	int		 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*a_    )






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



							'''simple docstring'''



							a__  :	List[str]		 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*a_    )



@require_tf
class 					__UpperCAmelCase							(						_UpperCamelCase						,  unittest.TestCase  ):
				__lowerCamelCase				:	Optional[Any]				     =				(
				    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
				)
				__lowerCamelCase				:	Any				     =				False
				__lowerCamelCase				:	Optional[int]				     =				False
				def       UpperCAmelCase      (		self    :			str    )	->	Optional[int]:



							'''simple docstring'''



							a__  :	Union[str, Any]		 =			TFFunnelModelTester(self	, base=a_    )
							a__  :	Optional[int]		 =			ConfigTester(self	, config_class=a_    )
				def       UpperCAmelCase      (		self    :			str    )	->	Union[str, Any]:



							'''simple docstring'''



							self.config_tester.run_common_tests()
				def       UpperCAmelCase      (		self    :			Union[str, Any]    )	->	Optional[Any]:



							'''simple docstring'''



							a__  :	Any		 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_base_model(*a_    )
				def       UpperCAmelCase      (		self    :			Tuple    )	->	List[Any]:



							'''simple docstring'''



							a__  :	Optional[Any]		 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*a_    )






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



							'''simple docstring'''



							a__  :	Optional[int]		 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*a_    )






"""simple docstring"""






import argparse
import pickle

import numpy as np
import torch
from torch import nn

from transformers import ReformerConfig, ReformerModelWithLMHead
from transformers.utils import logging


logging.set_verbosity_info()






def      lowercase__ (					lowerCAmelCase__   :			Union[str, Any]							,			lowerCAmelCase__   :			Dict							,			lowerCAmelCase__   :			List[Any]=None			)    ->			str:



			'''simple docstring'''
			# set parameter of one layer
			assert torch_layer.weight.shape == weight.shape, F"{torch_layer} layer.weight does not match"
			a__  :	Optional[Any]		 =			nn.Parameter(lowerCAmelCase__			)
			if bias is not None:
						assert torch_layer.bias.shape == bias.shape, F"{torch_layer} layer.bias does not match"
						a__  :	str		 =			nn.Parameter(lowerCAmelCase__			)






def      lowercase__ (					lowerCAmelCase__   :			Union[str, Any]							,			lowerCAmelCase__   :			Tuple							,			lowerCAmelCase__   :			Union[str, Any]			)    ->			str:



			'''simple docstring'''
			# set torch weights for 1-to-1 comparison
			a__  :	Tuple		 =			np.asarray(weights[0]			)
			a__  :	Any		 =			np.asarray(weights[1]			)
			a__  :	Optional[Any]		 =			np.asarray(weights[2]			)

			set_param(
			    torch_layer.self_attention.query_key							,			torch.tensor(lowerCAmelCase__			).transpose(1							,			2			).contiguous().view(-1							,			lowerCAmelCase__			)							,			)
			set_param(
			    torch_layer.self_attention.value							,			torch.tensor(lowerCAmelCase__			).transpose(1							,			2			).contiguous().view(-1							,			lowerCAmelCase__			)							,			)
			set_param(
			    torch_layer.output.dense							,			torch.tensor(lowerCAmelCase__			).view(-1							,			lowerCAmelCase__			).contiguous().transpose(0							,			1			)							,			)






def      lowercase__ (					lowerCAmelCase__   :			List[str]							,			lowerCAmelCase__   :			List[str]							,			lowerCAmelCase__   :			Union[str, Any]			)    ->			str:



			'''simple docstring'''
			# set torch weights for 1-to-1 comparison
			a__  :	int		 =			np.asarray(weights[0]			)
			a__  :	List[str]		 =			np.asarray(weights[1]			)
			a__  :	Optional[int]		 =			np.asarray(weights[2]			)
			a__  :	int		 =			np.asarray(weights[3]			)

			set_param(
			    torch_layer.self_attention.query							,			torch.tensor(lowerCAmelCase__			).transpose(1							,			2			).contiguous().view(-1							,			lowerCAmelCase__			)							,			)
			set_param(
			    torch_layer.self_attention.key							,			torch.tensor(lowerCAmelCase__			).transpose(1							,			2			).contiguous().view(-1							,			lowerCAmelCase__			)							,			)
			set_param(
			    torch_layer.self_attention.value							,			torch.tensor(lowerCAmelCase__			).transpose(1							,			2			).contiguous().view(-1							,			lowerCAmelCase__			)							,			)
			set_param(
			    torch_layer.output.dense							,			torch.tensor(lowerCAmelCase__			).view(-1							,			lowerCAmelCase__			).contiguous().transpose(0							,			1			)							,			)






def      lowercase__ (					lowerCAmelCase__   :			Any							,			lowerCAmelCase__   :			List[Any]							,			lowerCAmelCase__   :			Union[str, Any]			)    ->			Any:



			'''simple docstring'''
			# layernorm 1
			a__  :	int		 =			weights[0][0][0]
			a__  :	Optional[Any]		 =			np.asarray(layer_norm_a[0]			)
			a__  :	Optional[int]		 =			np.asarray(layer_norm_a[1]			)
			set_param(
			    torch_block.attention.layer_norm							,			torch.tensor(lowerCAmelCase__			)							,			torch.tensor(lowerCAmelCase__			)							,			)

			# lsh weights + output
			a__  :	Dict		 =			weights[0][1]
			if len(lowerCAmelCase__			) < 4:
						set_layer_weights_in_torch_lsh(lowerCAmelCase__							,			torch_block.attention							,			lowerCAmelCase__			)
			else:
						set_layer_weights_in_torch_local(lowerCAmelCase__							,			torch_block.attention							,			lowerCAmelCase__			)

			# intermediate weighs
			a__  :	Optional[int]		 =			weights[2][0][1][2]

			# Chunked Feed Forward
			if len(lowerCAmelCase__			) == 4:
						a__  :	List[str]		 =			intermediate_weights[2]

			# layernorm 2
			a__  :	List[Any]		 =			np.asarray(intermediate_weights[0][0]			)
			a__  :	Optional[int]		 =			np.asarray(intermediate_weights[0][1]			)
			set_param(
			    torch_block.feed_forward.layer_norm							,			torch.tensor(lowerCAmelCase__			)							,			torch.tensor(lowerCAmelCase__			)							,			)

			# intermediate dense
			a__  :	Tuple		 =			np.asarray(intermediate_weights[1][0]			)
			a__  :	Union[str, Any]		 =			np.asarray(intermediate_weights[1][1]			)
			set_param(
			    torch_block.feed_forward.dense.dense							,			torch.tensor(lowerCAmelCase__			).transpose(0							,			1			).contiguous()							,			torch.tensor(lowerCAmelCase__			)							,			)

			# intermediate out
			a__  :	Optional[Any]		 =			np.asarray(intermediate_weights[4][0]			)
			a__  :	Dict		 =			np.asarray(intermediate_weights[4][1]			)
			set_param(
			    torch_block.feed_forward.output.dense							,			torch.tensor(lowerCAmelCase__			).transpose(0							,			1			).contiguous()							,			torch.tensor(lowerCAmelCase__			)							,			)






def      lowercase__ (					lowerCAmelCase__   :			int							,			lowerCAmelCase__   :			List[str]							,			lowerCAmelCase__   :			Union[str, Any]			)    ->			Any:



			'''simple docstring'''
			# reformer model
			a__  :	int		 =			torch_model.reformer

			# word embeds
			a__  :	List[Any]		 =			np.asarray(weights[1]			)
			set_param(
			    torch_model_reformer.embeddings.word_embeddings							,			torch.tensor(lowerCAmelCase__			)							,			)

			if isinstance(weights[3]							,			lowerCAmelCase__			):
						a__  :	Dict		 =			torch_model_reformer.embeddings.position_embeddings
						for emb_idx in range(len(position_embeddings.weights			)			):
									a__  :	str		 =			np.asarray(weights[3][emb_idx][0]			)
									assert (
									    position_embeddings.weights[emb_idx].shape == emb_weights.shape
									), F"{position_embeddings[emb_idx]} emb does not match"
									a__  :	List[str]		 =			nn.Parameter(torch.tensor(lowerCAmelCase__			)			)

			a__  :	Optional[Any]		 =			weights[5]
			assert len(torch_model_reformer.encoder.layers			) * 4 == len(
			    lowerCAmelCase__			), "HF and trax model do not have the same number of layers"
			for layer_idx, layer in enumerate(torch_model_reformer.encoder.layers			):
						a__  :	Union[str, Any]		 =			trax_layer_weights[4 * layer_idx : 4 * (layer_idx + 1)]
						set_block_weights_in_torch(lowerCAmelCase__							,			lowerCAmelCase__							,			lowerCAmelCase__			)

			# output layer norm
			a__  :	str		 =			np.asarray(weights[7][0]			)
			a__  :	List[Any]		 =			np.asarray(weights[7][1]			)
			set_param(
			    torch_model_reformer.encoder.layer_norm							,			torch.tensor(lowerCAmelCase__			)							,			torch.tensor(lowerCAmelCase__			)							,			)

			# output embeddings
			a__  :	str		 =			np.asarray(weights[9][0]			)
			a__  :	Optional[int]		 =			np.asarray(weights[9][1]			)
			set_param(
			    torch_model.lm_head.decoder							,			torch.tensor(lowerCAmelCase__			).transpose(0							,			1			).contiguous()							,			torch.tensor(lowerCAmelCase__			)							,			)






def      lowercase__ (					lowerCAmelCase__   :			Tuple							,			lowerCAmelCase__   :			List[Any]							,			lowerCAmelCase__   :			Union[str, Any]			)    ->			List[str]:



			'''simple docstring'''
			# Initialise PyTorch model
			a__  :	List[Any]		 =			ReformerConfig.from_json_file(lowerCAmelCase__			)
			print(F"Building PyTorch model from configuration: {config}"			)
			a__  :	Optional[Any]		 =			ReformerModelWithLMHead(lowerCAmelCase__			)

			with open(lowerCAmelCase__							,			"rb"			) as f:
						a__  :	List[str]		 =			pickle.load(lowerCAmelCase__			)["weights"]

			set_model_weights_in_torch(lowerCAmelCase__							,			lowerCAmelCase__							,			config.hidden_size			)

			# Save pytorch-model
			print(F"Save PyTorch model to {pytorch_dump_path}"			)
			torch.save(model.state_dict()							,			lowerCAmelCase__			)


if __name__ == "__main__":
		__UpperCAmelCase						=			argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--trax_model_pkl_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
		)
		parser.add_argument(
		    '''--config_file''',
		    default=None,
		    type=str,
		    required=True,
		    help=(
		        '''The config json file corresponding to the pre-trained Reformer model. \n'''
		        '''This specifies the model architecture.'''
		    ),
		)
		parser.add_argument(
		    '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
		)
		__UpperCAmelCase						=			parser.parse_args()
		convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)