





"""simple docstring"""






import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


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

		from transformers.generation import (
		    FlaxForcedBOSTokenLogitsProcessor,
		    FlaxForcedEOSTokenLogitsProcessor,
		    FlaxLogitsProcessorList,
		    FlaxMinLengthLogitsProcessor,
		    FlaxTemperatureLogitsWarper,
		    FlaxTopKLogitsWarper,
		    FlaxTopPLogitsWarper,
		)



@require_flax
class 					__UpperCAmelCase							(						unittest.TestCase  ):
				def       UpperCAmelCase      (		self    :			Any	, a_    :			int	, a_    :			int    )	->	Union[str, Any]:



							'''simple docstring'''



							a__  :	int		 =			jnp.ones((batch_size, length)    ) / length
							return scores
				def       UpperCAmelCase      (		self    :			Any    )	->	Dict:



							'''simple docstring'''



							a__  :	str		 =			None
							a__  :	Optional[Any]		 =			20

							a__  :	str		 =			self._get_uniform_logits(batch_size=2	, length=a_    )

							# tweak scores to not be uniform anymore
							a__  :	str		 =			scores.at[1, 5].set((1 / length) + 0.1    )  # peak, 1st batch
							a__  :	Union[str, Any]		 =			scores.at[1, 10].set((1 / length) - 0.4    )  # valley, 1st batch

							# compute softmax
							a__  :	List[Any]		 =			jax.nn.softmax(a_	, axis=-1    )

							a__  :	Optional[Any]		 =			FlaxTemperatureLogitsWarper(temperature=0.5    )
							a__  :	List[Any]		 =			FlaxTemperatureLogitsWarper(temperature=1.3    )

							a__  :	str		 =			jax.nn.softmax(temp_dist_warper_sharper(a_	, scores.copy()	, cur_len=a_    )	, axis=-1    )
							a__  :	str		 =			jax.nn.softmax(temp_dist_warper_smoother(a_	, scores.copy()	, cur_len=a_    )	, axis=-1    )

							# uniform distribution stays uniform
							self.assertTrue(jnp.allclose(probs[0, :]	, warped_prob_sharp[0, :]	, atol=1E-3    )    )
							self.assertTrue(jnp.allclose(probs[0, :]	, warped_prob_smooth[0, :]	, atol=1E-3    )    )

							# sharp peaks get higher, valleys get lower
							self.assertLess(probs[1, :].max()	, warped_prob_sharp[1, :].max()    )
							self.assertGreater(probs[1, :].min()	, warped_prob_sharp[1, :].min()    )

							# smooth peaks get lower, valleys get higher
							self.assertGreater(probs[1, :].max()	, warped_prob_smooth[1, :].max()    )
							self.assertLess(probs[1, :].min()	, warped_prob_smooth[1, :].min()    )
				def       UpperCAmelCase      (		self    :			Optional[Any]    )	->	List[str]:



							'''simple docstring'''



							a__  :	List[Any]		 =			None
							a__  :	int		 =			10
							a__  :	List[str]		 =			2

							# create ramp distribution
							a__  :	str		 =			np.broadcast_to(np.arange(a_    )[None, :]	, (batch_size, vocab_size)    ).copy()
							a__  :	Tuple		 =			ramp_logits[1:, : vocab_size // 2] + vocab_size

							a__  :	Union[str, Any]		 =			FlaxTopKLogitsWarper(3    )

							a__  :	List[str]		 =			top_k_warp(a_	, a_	, cur_len=a_    )

							# check that correct tokens are filtered
							self.assertListEqual(jnp.isinf(scores[0]    ).tolist()	, 7 * [True] + 3 * [False]    )
							self.assertListEqual(jnp.isinf(scores[1]    ).tolist()	, 2 * [True] + 3 * [False] + 5 * [True]    )

							# check special case
							a__  :	Any		 =			5
							a__  :	Optional[int]		 =			FlaxTopKLogitsWarper(top_k=1	, filter_value=0.0	, min_tokens_to_keep=3    )

							a__  :	List[str]		 =			np.broadcast_to(np.arange(a_    )[None, :]	, (batch_size, length)    ).copy()
							a__  :	int		 =			top_k_warp_safety_check(a_	, a_	, cur_len=a_    )

							# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
							self.assertListEqual((scores == 0.0).sum(axis=-1    ).tolist()	, [2, 2]    )
				def       UpperCAmelCase      (		self    :			Optional[Any]    )	->	Any:



							'''simple docstring'''



							a__  :	Any		 =			None
							a__  :	List[str]		 =			10
							a__  :	int		 =			2

							# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
							a__  :	Optional[Any]		 =			np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]    )    )

							a__  :	Any		 =			FlaxTopPLogitsWarper(0.8    )
							a__  :	List[Any]		 =			np.exp(top_p_warp(a_	, a_	, cur_len=a_    )    )

							# dist should be filtered to keep min num values so that sum is >= top_p
							# exp (-inf) => 0
							a__  :	Tuple		 =			np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]    )
							self.assertTrue(np.allclose(a_	, a_	, atol=1E-3    )    )

							# check edge cases with negative and extreme logits
							a__  :	Union[str, Any]		 =			np.broadcast_to(np.arange(a_    )[None, :]	, (batch_size, vocab_size)    ).copy() - (
							    vocab_size // 2
							)

							# make ramp_logits more extreme
							a__  :	Dict		 =			ramp_logits[1] * 100.0

							# make sure at least 2 tokens are kept
							a__  :	List[str]		 =			FlaxTopPLogitsWarper(0.9	, min_tokens_to_keep=2	, filter_value=0.0    )
							a__  :	Any		 =			top_p_warp(a_	, a_	, cur_len=a_    )

							# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
							self.assertListEqual((filtered_dist != 0.0).sum(axis=-1    ).tolist()	, [3, 2]    )
				def       UpperCAmelCase      (		self    :			str    )	->	Optional[Any]:



							'''simple docstring'''



							a__  :	Optional[Any]		 =			20
							a__  :	str		 =			4
							a__  :	int		 =			0

							a__  :	Any		 =			FlaxMinLengthLogitsProcessor(min_length=10	, eos_token_id=a_    )

							# check that min length is applied at length 5
							a__  :	Optional[int]		 =			ids_tensor((batch_size, 20)	, vocab_size=20    )
							a__  :	Optional[Any]		 =			5
							a__  :	int		 =			self._get_uniform_logits(a_	, a_    )
							a__  :	str		 =			min_dist_processor(a_	, a_	, cur_len=a_    )
							self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()	, 4 * [-float("inf"    )]    )

							# check that min length is not applied anymore at length 15
							a__  :	List[Any]		 =			self._get_uniform_logits(a_	, a_    )
							a__  :	Optional[int]		 =			15
							a__  :	Optional[int]		 =			min_dist_processor(a_	, a_	, cur_len=a_    )
							self.assertFalse(jnp.isinf(a_    ).any()    )
				def       UpperCAmelCase      (		self    :			Union[str, Any]    )	->	Any:



							'''simple docstring'''



							a__  :	List[Any]		 =			20
							a__  :	Tuple		 =			4
							a__  :	Optional[int]		 =			0

							a__  :	Optional[int]		 =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_    )

							# check that all scores are -inf except the bos_token_id score
							a__  :	Any		 =			ids_tensor((batch_size, 1)	, vocab_size=20    )
							a__  :	Union[str, Any]		 =			1
							a__  :	Dict		 =			self._get_uniform_logits(a_	, a_    )
							a__  :	List[str]		 =			logits_processor(a_	, a_	, cur_len=a_    )
							self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]    ).all()    )
							self.assertListEqual(scores[:, bos_token_id].tolist()	, 4 * [0]    )  # score for bos_token_id shold be zero

							# check that bos_token_id is not forced if current length is greater than 1
							a__  :	str		 =			3
							a__  :	Optional[int]		 =			self._get_uniform_logits(a_	, a_    )
							a__  :	Optional[int]		 =			logits_processor(a_	, a_	, cur_len=a_    )
							self.assertFalse(jnp.isinf(a_    ).any()    )
				def       UpperCAmelCase      (		self    :			str    )	->	List[Any]:



							'''simple docstring'''



							a__  :	Union[str, Any]		 =			20
							a__  :	Tuple		 =			4
							a__  :	List[Any]		 =			0
							a__  :	Optional[Any]		 =			5

							a__  :	int		 =			FlaxForcedEOSTokenLogitsProcessor(max_length=a_	, eos_token_id=a_    )

							# check that all scores are -inf except the eos_token_id when max_length is reached
							a__  :	Tuple		 =			ids_tensor((batch_size, 4)	, vocab_size=20    )
							a__  :	Tuple		 =			4
							a__  :	List[str]		 =			self._get_uniform_logits(a_	, a_    )
							a__  :	List[Any]		 =			logits_processor(a_	, a_	, cur_len=a_    )
							self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]    ).all()    )
							self.assertListEqual(scores[:, eos_token_id].tolist()	, 4 * [0]    )  # score for eos_token_id should be zero

							# check that eos_token_id is not forced if max_length is not reached
							a__  :	List[Any]		 =			3
							a__  :	Optional[int]		 =			self._get_uniform_logits(a_	, a_    )
							a__  :	Tuple		 =			logits_processor(a_	, a_	, cur_len=a_    )
							self.assertFalse(jnp.isinf(a_    ).any()    )
				def       UpperCAmelCase      (		self    :			List[str]    )	->	Tuple:



							'''simple docstring'''



							a__  :	Optional[int]		 =			4
							a__  :	List[Any]		 =			10
							a__  :	Any		 =			15
							a__  :	Optional[int]		 =			2
							a__  :	List[str]		 =			1
							a__  :	Optional[Any]		 =			15

							# dummy input_ids and scores
							a__  :	Union[str, Any]		 =			ids_tensor((batch_size, sequence_length)	, a_    )
							a__  :	List[Any]		 =			input_ids.copy()

							a__  :	List[str]		 =			self._get_uniform_logits(a_	, a_    )
							a__  :	int		 =			scores.copy()

							# instantiate all dist processors
							a__  :	Any		 =			FlaxTemperatureLogitsWarper(temperature=0.5    )
							a__  :	Any		 =			FlaxTopKLogitsWarper(3    )
							a__  :	Optional[int]		 =			FlaxTopPLogitsWarper(0.8    )

							# instantiate all logits processors
							a__  :	Tuple		 =			FlaxMinLengthLogitsProcessor(min_length=10	, eos_token_id=a_    )
							a__  :	List[str]		 =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_    )
							a__  :	int		 =			FlaxForcedEOSTokenLogitsProcessor(max_length=a_	, eos_token_id=a_    )

							a__  :	str		 =			10

							# no processor list
							a__  :	Tuple		 =			temp_dist_warp(a_	, a_	, cur_len=a_    )
							a__  :	Dict		 =			top_k_warp(a_	, a_	, cur_len=a_    )
							a__  :	Optional[Any]		 =			top_p_warp(a_	, a_	, cur_len=a_    )
							a__  :	Tuple		 =			min_dist_proc(a_	, a_	, cur_len=a_    )
							a__  :	Dict		 =			bos_dist_proc(a_	, a_	, cur_len=a_    )
							a__  :	Union[str, Any]		 =			eos_dist_proc(a_	, a_	, cur_len=a_    )

							# with processor list
							a__  :	str		 =			FlaxLogitsProcessorList(
							    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]    )
							a__  :	List[Any]		 =			processor(a_	, a_	, cur_len=a_    )

							# scores should be equal
							self.assertTrue(jnp.allclose(a_	, a_	, atol=1E-3    )    )

							# input_ids should never be changed
							self.assertListEqual(input_ids.tolist()	, input_ids_comp.tolist()    )






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



							'''simple docstring'''



							a__  :	Optional[int]		 =			4
							a__  :	Optional[Any]		 =			10
							a__  :	str		 =			15
							a__  :	Any		 =			2
							a__  :	Dict		 =			1
							a__  :	Optional[Any]		 =			15

							# dummy input_ids and scores
							a__  :	Any		 =			ids_tensor((batch_size, sequence_length)	, a_    )
							a__  :	Tuple		 =			input_ids.copy()

							a__  :	Dict		 =			self._get_uniform_logits(a_	, a_    )
							a__  :	Optional[int]		 =			scores.copy()

							# instantiate all dist processors
							a__  :	Union[str, Any]		 =			FlaxTemperatureLogitsWarper(temperature=0.5    )
							a__  :	List[Any]		 =			FlaxTopKLogitsWarper(3    )
							a__  :	str		 =			FlaxTopPLogitsWarper(0.8    )

							# instantiate all logits processors
							a__  :	int		 =			FlaxMinLengthLogitsProcessor(min_length=10	, eos_token_id=a_    )
							a__  :	Optional[int]		 =			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_    )
							a__  :	Dict		 =			FlaxForcedEOSTokenLogitsProcessor(max_length=a_	, eos_token_id=a_    )

							a__  :	Optional[int]		 =			10

							# no processor list
							def run_no_processor_list(a_    :			Optional[Any]	, a_    :			List[str]	, a_    :			List[Any]    ):
										a__  :	Union[str, Any]		 =			temp_dist_warp(a_	, a_	, cur_len=a_    )
										a__  :	List[str]		 =			top_k_warp(a_	, a_	, cur_len=a_    )
										a__  :	List[str]		 =			top_p_warp(a_	, a_	, cur_len=a_    )
										a__  :	Optional[Any]		 =			min_dist_proc(a_	, a_	, cur_len=a_    )
										a__  :	List[str]		 =			bos_dist_proc(a_	, a_	, cur_len=a_    )
										a__  :	Dict		 =			eos_dist_proc(a_	, a_	, cur_len=a_    )
										return scores

							# with processor list
							def run_processor_list(a_    :			List[str]	, a_    :			Dict	, a_    :			Dict    ):
										a__  :	str		 =			FlaxLogitsProcessorList(
										    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]    )
										a__  :	int		 =			processor(a_	, a_	, cur_len=a_    )
										return scores

							a__  :	Any		 =			jax.jit(a_    )
							a__  :	List[Any]		 =			jax.jit(a_    )

							a__  :	Optional[Any]		 =			jitted_run_no_processor_list(a_	, a_	, a_    )
							a__  :	List[Any]		 =			jitted_run_processor_list(a_	, a_	, a_    )

							# scores should be equal
							self.assertTrue(jnp.allclose(a_	, a_	, atol=1E-3    )    )

							# input_ids should never be changed
							self.assertListEqual(input_ids.tolist()	, input_ids_comp.tolist()    )






"""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)