




'''simple docstring'''







import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       MODEL_FOR_PRETRAINING_MAPPING,
       MegatronBertForCausalLM,
       MegatronBertForMaskedLM,
       MegatronBertForMultipleChoice,
       MegatronBertForNextSentencePrediction,
       MegatronBertForPreTraining,
       MegatronBertForQuestionAnswering,
       MegatronBertForSequenceClassification,
       MegatronBertForTokenClassification,
       MegatronBertModel,
   )


class 			_a  :
       def __init__(	self						:	Optional[Any]   ,				lowercase						:	List[str]   ,				lowercase						:	str=13   ,				lowercase						:	List[Any]=7   ,				lowercase						:	List[str]=True   ,				lowercase						:	Optional[int]=True   ,				lowercase						:	List[str]=True   ,				lowercase						:	Dict=True   ,				lowercase						:	Optional[Any]=99   ,				lowercase						:	Any=64   ,				lowercase						:	Optional[int]=32   ,				lowercase						:	List[str]=5   ,				lowercase						:	str=4   ,				lowercase						:	Union[str, Any]=37   ,				lowercase						:	Union[str, Any]="gelu"   ,				lowercase						:	Dict=0.1   ,				lowercase						:	Optional[Any]=0.1   ,				lowercase						:	Dict=512   ,				lowercase						:	Union[str, Any]=16   ,				lowercase						:	Optional[int]=2   ,				lowercase						:	str=0.02   ,				lowercase						:	Dict=3   ,				lowercase						:	str=4   ,				lowercase						:	Tuple=None   ,				):



           '''simple docstring'''

           UpperCAmelCase			       =		parent
           UpperCAmelCase			       =		batch_size
           UpperCAmelCase			       =		seq_length
           UpperCAmelCase			       =		is_training
           UpperCAmelCase			       =		use_input_mask
           UpperCAmelCase			       =		use_token_type_ids
           UpperCAmelCase			       =		use_labels
           UpperCAmelCase			       =		vocab_size
           UpperCAmelCase			       =		hidden_size
           UpperCAmelCase			       =		embedding_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			       =		max_position_embeddings
           UpperCAmelCase			       =		type_vocab_size
           UpperCAmelCase			       =		type_sequence_label_size
           UpperCAmelCase			       =		initializer_range
           UpperCAmelCase			       =		num_labels
           UpperCAmelCase			       =		num_choices
           UpperCAmelCase			       =		scope
       def     A      (	self						:	Optional[int]				):



           '''simple docstring'''

           UpperCAmelCase			       =		ids_tensor([self.batch_size, self.seq_length]   ,				self.vocab_size				)

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

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

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

           UpperCAmelCase			       =		self.get_config()

           return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
       def     A      (	self						:	List[Any]				):



           '''simple docstring'''

           return MegatronBertConfig(
               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   ,				embedding_size=self.embedding_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=_a   ,				initializer_range=self.initializer_range   ,				)
       def     A      (	self						:	Tuple   ,				lowercase						:	int   ,				lowercase						:	str   ,				lowercase						:	str   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	str   ,				lowercase						:	int   ,				lowercase						:	Any				):



           '''simple docstring'''

           UpperCAmelCase			       =		MegatronBertModel(config=_a				)
           model.to(_a				)
           model.eval()
           UpperCAmelCase			       =		model(_a   ,				attention_mask=_a   ,				token_type_ids=_a				)
           UpperCAmelCase			       =		model(_a   ,				token_type_ids=_a				)
           UpperCAmelCase			       =		model(_a				)

           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						:	Union[str, Any]   ,				lowercase						:	Dict   ,				lowercase						:	str   ,				lowercase						:	int   ,				lowercase						:	str   ,				lowercase						:	Dict   ,				lowercase						:	Optional[Any]   ,				lowercase						:	str				):



           '''simple docstring'''

           UpperCAmelCase			       =		MegatronBertForMaskedLM(config=_a				)
           model.to(_a				)
           model.eval()
           UpperCAmelCase			       =		model(_a   ,				attention_mask=_a   ,				token_type_ids=_a   ,				labels=_a				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length, self.vocab_size)				)
       def     A      (	self						:	Optional[Any]   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	Any   ,				lowercase						:	Optional[Any]   ,				lowercase						:	int   ,				lowercase						:	str   ,				lowercase						:	str   ,				lowercase						:	List[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		MegatronBertForCausalLM(config=_a				)
           model.to(_a				)
           model.eval()
           UpperCAmelCase			       =		model(_a   ,				attention_mask=_a   ,				token_type_ids=_a   ,				labels=_a				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length, self.vocab_size)				)
       def     A      (	self						:	Any   ,				lowercase						:	List[str]   ,				lowercase						:	List[str]   ,				lowercase						:	List[str]   ,				lowercase						:	List[str]   ,				lowercase						:	List[str]   ,				lowercase						:	List[Any]   ,				lowercase						:	Tuple				):



           '''simple docstring'''

           UpperCAmelCase			       =		MegatronBertForNextSentencePrediction(config=_a				)
           model.to(_a				)
           model.eval()
           UpperCAmelCase			       =		model(
               _a   ,				attention_mask=_a   ,				token_type_ids=_a   ,				labels=_a   ,				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, 2)				)
       def     A      (	self						:	Tuple   ,				lowercase						:	List[str]   ,				lowercase						:	List[Any]   ,				lowercase						:	Any   ,				lowercase						:	List[str]   ,				lowercase						:	str   ,				lowercase						:	List[str]   ,				lowercase						:	Any				):



           '''simple docstring'''

           UpperCAmelCase			       =		MegatronBertForPreTraining(config=_a				)
           model.to(_a				)
           model.eval()
           UpperCAmelCase			       =		model(
               _a   ,				attention_mask=_a   ,				token_type_ids=_a   ,				labels=_a   ,				next_sentence_label=_a   ,				)
           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						:	str   ,				lowercase						:	str   ,				lowercase						:	int   ,				lowercase						:	str   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	Optional[int]   ,				lowercase						:	Dict				):



           '''simple docstring'''

           UpperCAmelCase			       =		MegatronBertForQuestionAnswering(config=_a				)
           model.to(_a				)
           model.eval()
           UpperCAmelCase			       =		model(
               _a   ,				attention_mask=_a   ,				token_type_ids=_a   ,				start_positions=_a   ,				end_positions=_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     A      (	self						:	str   ,				lowercase						:	int   ,				lowercase						:	Dict   ,				lowercase						:	List[str]   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	Tuple   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	List[str]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.num_labels
           UpperCAmelCase			       =		MegatronBertForSequenceClassification(_a				)
           model.to(_a				)
           model.eval()
           UpperCAmelCase			       =		model(_a   ,				attention_mask=_a   ,				token_type_ids=_a   ,				labels=_a				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_labels)				)
       def     A      (	self						:	Tuple   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	Dict   ,				lowercase						:	List[str]   ,				lowercase						:	Tuple   ,				lowercase						:	List[str]   ,				lowercase						:	Any   ,				lowercase						:	Optional[int]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.num_labels
           UpperCAmelCase			       =		MegatronBertForTokenClassification(config=_a				)
           model.to(_a				)
           model.eval()
           UpperCAmelCase			       =		model(_a   ,				attention_mask=_a   ,				token_type_ids=_a   ,				labels=_a				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length, self.num_labels)				)
       def     A      (	self						:	Any   ,				lowercase						:	Optional[int]   ,				lowercase						:	Optional[Any]   ,				lowercase						:	Dict   ,				lowercase						:	int   ,				lowercase						:	str   ,				lowercase						:	List[str]   ,				lowercase						:	Dict				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.num_choices
           UpperCAmelCase			       =		MegatronBertForMultipleChoice(config=_a				)
           model.to(_a				)
           model.eval()
           UpperCAmelCase			       =		input_ids.unsqueeze(1				).expand(-1   ,				self.num_choices   ,				-1				).contiguous()
           UpperCAmelCase			       =		token_type_ids.unsqueeze(1				).expand(-1   ,				self.num_choices   ,				-1				).contiguous()
           UpperCAmelCase			       =		input_mask.unsqueeze(1				).expand(-1   ,				self.num_choices   ,				-1				).contiguous()
           UpperCAmelCase			       =		model(
               _a   ,				attention_mask=_a   ,				token_type_ids=_a   ,				labels=_a   ,				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_choices)				)







       def     A      (	self						:	Dict				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.prepare_config_and_inputs()
           (
               UpperCAmelCase  
           )			       =		config_and_inputs
           UpperCAmelCase			       =		{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
           return config, inputs_dict






@require_torch
class 			_a  (	lowercase__							,							lowercase__							,							unittest.TestCase ):
       __a				:     str   				=					(
           (
               MegatronBertModel,
               MegatronBertForMaskedLM,
               MegatronBertForCausalLM,
               MegatronBertForMultipleChoice,
               MegatronBertForNextSentencePrediction,
               MegatronBertForPreTraining,
               MegatronBertForQuestionAnswering,
               MegatronBertForSequenceClassification,
               MegatronBertForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       __a				:     Tuple   				=					(
           {
               '''feature-extraction''': MegatronBertModel,
               '''fill-mask''': MegatronBertForMaskedLM,
               '''question-answering''': MegatronBertForQuestionAnswering,
               '''text-classification''': MegatronBertForSequenceClassification,
               '''text-generation''': MegatronBertForCausalLM,
               '''token-classification''': MegatronBertForTokenClassification,
               '''zero-shot''': MegatronBertForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       __a				:     List[Any]   				=					True
       # test_resize_embeddings = False
       __a				:     List[Any]   				=					False
       def     A      (	self						:	Any   ,				lowercase						:	List[Any]   ,				lowercase						:	List[str]   ,				lowercase						:	List[str]=False				):



           '''simple docstring'''

           UpperCAmelCase			       =		super()._prepare_for_class(_a   ,				_a   ,				return_labels=_a				)

           if return_labels:
               if model_class in get_values(_a				):
                   UpperCAmelCase			       =		torch.zeros(
                       (self.model_tester.batch_size, self.model_tester.seq_length)   ,				dtype=torch.long   ,				device=_a				)
                   UpperCAmelCase			       =		torch.zeros(
                       self.model_tester.batch_size   ,				dtype=torch.long   ,				device=_a				)
           return inputs_dict
       def     A      (	self						:	Union[str, Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		MegatronBertModelTester(self				)
           UpperCAmelCase			       =		ConfigTester(self   ,				config_class=_a   ,				hidden_size=37				)
       def     A      (	self						:	str				):



           '''simple docstring'''

           self.config_tester.run_common_tests()
       def     A      (	self						:	Optional[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_model(*_a				)
       def     A      (	self						:	Optional[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_a				)
       def     A      (	self						:	int				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_a				)
       def     A      (	self						:	Union[str, Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_a				)
       def     A      (	self						:	List[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_pretraining(*_a				)
       def     A      (	self						:	Any				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_question_answering(*_a				)
       def     A      (	self						:	Any				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_a				)







       def     A      (	self						:	List[str]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_token_classification(*_a				)






def 						snake_case_							(_a	:							Tuple     ):
    return torch.tensor(
        __a    ,   dtype=torch.long    ,   device=__a    ,   )


A     =1E-4







@require_torch
@require_sentencepiece
@require_tokenizers
class 			_a  (	unittest.TestCase ):
       @slow
       @unittest.skip('''Model is not available.'''				)
       def     A      (	self						:	str				):



           '''simple docstring'''

           UpperCAmelCase			       =		'nvidia/megatron-bert-uncased-345m'
           if "MYDIR" in os.environ:
               UpperCAmelCase			       =		os.path.join(os.environ['''MYDIR''']   ,				_a				)
           UpperCAmelCase			       =		MegatronBertModel.from_pretrained(_a				)
           model.to(_a				)
           model.half()
           UpperCAmelCase			       =		_long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]				)
           with torch.no_grad():
               UpperCAmelCase			       =		model(_a				)[0]
           UpperCAmelCase			       =		torch.Size((1, 9, 1_024)				)
           self.assertEqual(output.shape   ,				_a				)

           UpperCAmelCase			       =		[-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
           for ii in range(3				):
               for jj in range(3				):
                   UpperCAmelCase			       =		output[0, ii, jj]
                   UpperCAmelCase			       =		expected[3 * ii + jj]
                   UpperCAmelCase			       =		'ii={} jj={} a={} b={}'.format(_a   ,				_a   ,				_a   ,				_a				)
                   self.assertTrue(math.isclose(_a   ,				_a   ,				rel_tol=_a   ,				abs_tol=_a				)   ,				msg=_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 					UpperCAmelCase_		(__a      :    Optional[Any] ,	__a      :    str ,	__a      :    Optional[Any]=None ):






   """simple docstring"""

   assert torch_layer.weight.shape == weight.shape, f"""{torch_layer} layer.weight does not match"""
   _a				:	str			=   nn.Parameter(__a )
   if bias is not None:
      assert torch_layer.bias.shape == bias.shape, f"""{torch_layer} layer.bias does not match"""
      _a				:	Any			=   nn.Parameter(__a )





def 					UpperCAmelCase_		(__a      :    int ,	__a      :    Optional[Any] ,	__a      :    int ):






   """simple docstring"""

   _a				:	Tuple			=   np.asarray(weights[0] )
   _a				:	Union[str, Any]			=   np.asarray(weights[1] )
   _a				:	Dict			=   np.asarray(weights[2] )

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





def 					UpperCAmelCase_		(__a      :    Optional[Any] ,	__a      :    Optional[int] ,	__a      :    List[str] ):






   """simple docstring"""

   _a				:	Dict			=   np.asarray(weights[0] )
   _a				:	Union[str, Any]			=   np.asarray(weights[1] )
   _a				:	str			=   np.asarray(weights[2] )
   _a				:	int			=   np.asarray(weights[3] )

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





def 					UpperCAmelCase_		(__a      :    Any ,	__a      :    Any ,	__a      :    Optional[Any] ):






   """simple docstring"""

   _a				:	List[str]			=   weights[0][0][0]
   _a				:	List[Any]			=   np.asarray(layer_norm_a[0] )
   _a				:	List[str]			=   np.asarray(layer_norm_a[1] )
   set_param(
       torch_block.attention.layer_norm ,	torch.tensor(__a ) ,	torch.tensor(__a ) ,	)

   # lsh weights + output
   _a				:	List[str]			=   weights[0][1]
   if len(__a ) < 4:
      set_layer_weights_in_torch_lsh(__a ,	torch_block.attention ,	__a )
   else:
      set_layer_weights_in_torch_local(__a ,	torch_block.attention ,	__a )

   # intermediate weighs
   _a				:	Optional[Any]			=   weights[2][0][1][2]

   # Chunked Feed Forward
   if len(__a ) == 4:
      _a				:	Union[str, Any]			=   intermediate_weights[2]

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

   # intermediate dense
   _a				:	Any			=   np.asarray(intermediate_weights[1][0] )
   _a				:	Any			=   np.asarray(intermediate_weights[1][1] )
   set_param(
       torch_block.feed_forward.dense.dense ,	torch.tensor(__a ).transpose(0 ,	1 ).contiguous() ,	torch.tensor(__a ) ,	)

   # intermediate out
   _a				:	Optional[int]			=   np.asarray(intermediate_weights[4][0] )
   _a				:	int			=   np.asarray(intermediate_weights[4][1] )
   set_param(
       torch_block.feed_forward.output.dense ,	torch.tensor(__a ).transpose(0 ,	1 ).contiguous() ,	torch.tensor(__a ) ,	)





def 					UpperCAmelCase_		(__a      :    Dict ,	__a      :    Dict ,	__a      :    List[Any] ):






   """simple docstring"""

   _a				:	Optional[int]			=   torch_model.reformer

   # word embeds
   _a				:	Tuple			=   np.asarray(weights[1] )
   set_param(
       torch_model_reformer.embeddings.word_embeddings ,	torch.tensor(__a ) ,	)

   if isinstance(weights[3] ,	__a ):
      _a				:	Any			=   torch_model_reformer.embeddings.position_embeddings
      for emb_idx in range(len(position_embeddings.weights ) ):
         _a				:	List[Any]			=   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				:	Any			=   nn.Parameter(torch.tensor(__a ) )

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

   # output layer norm
   _a				:	Optional[Any]			=   np.asarray(weights[7][0] )
   _a				:	int			=   np.asarray(weights[7][1] )
   set_param(
       torch_model_reformer.encoder.layer_norm ,	torch.tensor(__a ) ,	torch.tensor(__a ) ,	)

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





def 					UpperCAmelCase_		(__a      :    Tuple ,	__a      :    Optional[Any] ,	__a      :    Dict ):






   """simple docstring"""

   _a				:	List[Any]			=   ReformerConfig.from_json_file(__a )
   print(f"""Building PyTorch model from configuration: {config}""" )
   _a				:	int			=   ReformerModelWithLMHead(__a )

   with open(__a ,	'rb' ) as f:
      _a				:	Optional[Any]			=   pickle.load(__a )['weights']

   set_model_weights_in_torch(__a ,	__a ,	config.hidden_size )

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


if __name__ == "__main__":
  __lowerCAmelCase         =    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."""
  )
  __lowerCAmelCase         =    parser.parse_args()
  convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)



