





from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


__A				:    str       =   logging.get_logger(__name__)

# General docstring
__A				:    Dict       =   "ResNetConfig"

# Base docstring
__A				:    int       =   "microsoft/resnet-50"
__A				:    Optional[Any]       =   [1, 2_048, 7, 7]

# Image classification docstring
__A				:    Any       =   "microsoft/resnet-50"
__A				:    str       =   "tiger cat"

__A				:    int       =   [
    "microsoft/resnet-50",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]

class 	A_       (nn.Module			):
	def __init__(    self							,					_A							,					_A							,					_A = 3							,					_A = 1							,					_A = "relu"		):






				'''simple docstring'''







				super().__init__()
				UpperCAmelCase   		=				nn.Convad(
				    _A							,					_A							,					kernel_size=_A							,					stride=_A							,					padding=kernel_size // 2							,					bias=_A		)
				UpperCAmelCase   		=				nn.BatchNormad(_A		)
				UpperCAmelCase   		=				ACTaFN[activation] if activation is not None else nn.Identity()







	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.convolution(_A		)
				UpperCAmelCase   		=				self.normalization(_A		)
				UpperCAmelCase   		=				self.activation(_A		)
				return hidden_state

class 	A_       (nn.Module			):
	def __init__(    self							,					_A		):






				'''simple docstring'''







				super().__init__()
				UpperCAmelCase   		=				ResNetConvLayer(
				    config.num_channels							,					config.embedding_size							,					kernel_size=7							,					stride=2							,					activation=config.hidden_act		)
				UpperCAmelCase   		=				nn.MaxPoolad(kernel_size=3							,					stride=2							,					padding=1		)
				UpperCAmelCase   		=				config.num_channels







	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				UpperCAmelCase   		=				pixel_values.shape[1]
				if num_channels != self.num_channels:
							raise ValueError(
							    '''Make sure that the channel dimension of the pixel values match with the one set in the configuration.'''		)
				UpperCAmelCase   		=				self.embedder(_A		)
				UpperCAmelCase   		=				self.pooler(_A		)
				return embedding

class 	A_       (nn.Module			):
	def __init__(    self							,					_A							,					_A							,					_A = 2		):






				'''simple docstring'''







				super().__init__()
				UpperCAmelCase   		=				nn.Convad(_A							,					_A							,					kernel_size=1							,					stride=_A							,					bias=_A		)
				UpperCAmelCase   		=				nn.BatchNormad(_A		)







	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.convolution(_A		)
				UpperCAmelCase   		=				self.normalization(_A		)
				return hidden_state

class 	A_       (nn.Module			):
	def __init__(    self							,					_A							,					_A							,					_A = 1							,					_A = "relu"		):






				'''simple docstring'''







				super().__init__()
				UpperCAmelCase   		=				in_channels != out_channels or stride != 1
				UpperCAmelCase   		=				(
				    ResNetShortCut(_A							,					_A							,					stride=_A		) if should_apply_shortcut else nn.Identity()
				)
				UpperCAmelCase   		=				nn.Sequential(
				    ResNetConvLayer(_A							,					_A							,					stride=_A		)							,					ResNetConvLayer(_A							,					_A							,					activation=_A		)							,					)
				UpperCAmelCase   		=				ACTaFN[activation]







	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				UpperCAmelCase   		=				hidden_state
				UpperCAmelCase   		=				self.layer(_A		)
				UpperCAmelCase   		=				self.shortcut(_A		)
				hidden_state += residual
				UpperCAmelCase   		=				self.activation(_A		)
				return hidden_state

class 	A_       (nn.Module			):
	def __init__(    self							,					_A							,					_A							,					_A = 1							,					_A = "relu"							,					_A = 4		):






				'''simple docstring'''







				super().__init__()
				UpperCAmelCase   		=				in_channels != out_channels or stride != 1
				UpperCAmelCase   		=				out_channels // reduction
				UpperCAmelCase   		=				(
				    ResNetShortCut(_A							,					_A							,					stride=_A		) if should_apply_shortcut else nn.Identity()
				)
				UpperCAmelCase   		=				nn.Sequential(
				    ResNetConvLayer(_A							,					_A							,					kernel_size=1		)							,					ResNetConvLayer(_A							,					_A							,					stride=_A		)							,					ResNetConvLayer(_A							,					_A							,					kernel_size=1							,					activation=_A		)							,					)
				UpperCAmelCase   		=				ACTaFN[activation]







	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				UpperCAmelCase   		=				hidden_state
				UpperCAmelCase   		=				self.layer(_A		)
				UpperCAmelCase   		=				self.shortcut(_A		)
				hidden_state += residual
				UpperCAmelCase   		=				self.activation(_A		)
				return hidden_state

class 	A_       (nn.Module			):
	def __init__(    self							,					_A							,					_A							,					_A							,					_A = 2							,					_A = 2							,					):






				'''simple docstring'''







				super().__init__()

				UpperCAmelCase   		=				ResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else ResNetBasicLayer

				UpperCAmelCase   		=				nn.Sequential(
				    # downsampling is done in the first layer with stride of 2
				    layer(_A							,					_A							,					stride=_A							,					activation=config.hidden_act		)							,					*[layer(_A							,					_A							,					activation=config.hidden_act		) for _ in range(depth - 1		)]							,					)







	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				UpperCAmelCase   		=				input
				for layer in self.layers:
							UpperCAmelCase   		=				layer(_A		)
				return hidden_state

class 	A_       (nn.Module			):
	def __init__(    self							,					_A		):






				'''simple docstring'''







				super().__init__()
				UpperCAmelCase   		=				nn.ModuleList([]		)
				# based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
				self.stages.append(
				    ResNetStage(
				        _A							,					config.embedding_size							,					config.hidden_sizes[0]							,					stride=2 if config.downsample_in_first_stage else 1							,					depth=config.depths[0]							,					)		)
				UpperCAmelCase   		=				zip(config.hidden_sizes							,					config.hidden_sizes[1:]		)
				for (in_channels, out_channels), depth in zip(_A							,					config.depths[1:]		):
							self.stages.append(ResNetStage(_A							,					_A							,					_A							,					depth=_A		)		)







	def      _lowercase				(    self							,					_A							,					_A = False							,					_A = True		):






				'''simple docstring'''







				UpperCAmelCase   		=				() if output_hidden_states else None

				for stage_module in self.stages:
							if output_hidden_states:
										UpperCAmelCase   		=				hidden_states + (hidden_state,)

							UpperCAmelCase   		=				stage_module(_A		)

				if output_hidden_states:
							UpperCAmelCase   		=				hidden_states + (hidden_state,)

				if not return_dict:
							return tuple(v for v in [hidden_state, hidden_states] if v is not None		)

				return BaseModelOutputWithNoAttention(
				    last_hidden_state=_A							,					hidden_states=_A							,					)

class 	A_       (a_			):
	UpperCAmelCase__					=			ResNetConfig
	UpperCAmelCase__					=			'''resnet'''
	UpperCAmelCase__					=			'''pixel_values'''
	UpperCAmelCase__					=			True
	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				if isinstance(_A							,					nn.Convad		):
							nn.init.kaiming_normal_(module.weight							,					mode='''fan_out'''							,					nonlinearity='''relu'''		)
				elif isinstance(_A							,					(nn.BatchNormad, nn.GroupNorm)		):
							nn.init.constant_(module.weight							,					1		)
							nn.init.constant_(module.bias							,					0		)







	def      _lowercase				(    self							,					_A							,					_A=False		):






				'''simple docstring'''







				if isinstance(_A							,					_A		):
							UpperCAmelCase   		=				value
__A				:    int       =   R"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

__A				:    Tuple       =   R"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"

@add_start_docstrings(
    '''The bare ResNet model outputting raw features without any specific head on top.'''     ,							a_     ,							)
class 	A_       (a_			):
	def __init__(    self							,					_A		):






				'''simple docstring'''







				super().__init__(_A		)
				UpperCAmelCase   		=				config
				UpperCAmelCase   		=				ResNetEmbeddings(_A		)
				UpperCAmelCase   		=				ResNetEncoder(_A		)
				UpperCAmelCase   		=				nn.AdaptiveAvgPoolad((1, 1)		)
				# Initialize weights and apply final processing
				self.post_init()







	@add_start_docstrings_to_model_forward(_A		)
	@add_code_sample_docstrings(
	    checkpoint=_CHECKPOINT_FOR_DOC							,					output_type=_A							,					config_class=_CONFIG_FOR_DOC							,					modality='''vision'''							,					expected_output=_EXPECTED_OUTPUT_SHAPE							,					)
	def      _lowercase				(    self							,					_A							,					_A = None							,					_A = None		):






				'''simple docstring'''







				UpperCAmelCase   		=				(
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				UpperCAmelCase   		=				return_dict if return_dict is not None else self.config.use_return_dict

				UpperCAmelCase   		=				self.embedder(_A		)

				UpperCAmelCase   		=				self.encoder(
				    _A							,					output_hidden_states=_A							,					return_dict=_A		)

				UpperCAmelCase   		=				encoder_outputs[0]

				UpperCAmelCase   		=				self.pooler(_A		)

				if not return_dict:
							return (last_hidden_state, pooled_output) + encoder_outputs[1:]

				return BaseModelOutputWithPoolingAndNoAttention(
				    last_hidden_state=_A							,					pooler_output=_A							,					hidden_states=encoder_outputs.hidden_states							,					)

@add_start_docstrings(
    '''
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    '''     ,							a_     ,							)
class 	A_       (a_			):
	def __init__(    self							,					_A		):






				'''simple docstring'''







				super().__init__(_A		)
				UpperCAmelCase   		=				config.num_labels
				UpperCAmelCase   		=				ResNetModel(_A		)
				# classification head
				UpperCAmelCase   		=				nn.Sequential(
				    nn.Flatten()							,					nn.Linear(config.hidden_sizes[-1]							,					config.num_labels		) if config.num_labels > 0 else nn.Identity()							,					)
				# initialize weights and apply final processing
				self.post_init()







	@add_start_docstrings_to_model_forward(_A		)
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT							,					output_type=_A							,					config_class=_CONFIG_FOR_DOC							,					expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							,					)
	def      _lowercase				(    self							,					_A = None							,					_A = None							,					_A = None							,					_A = None							,					):






				'''simple docstring'''







				UpperCAmelCase   		=				return_dict if return_dict is not None else self.config.use_return_dict

				UpperCAmelCase   		=				self.resnet(_A							,					output_hidden_states=_A							,					return_dict=_A		)

				UpperCAmelCase   		=				outputs.pooler_output if return_dict else outputs[1]

				UpperCAmelCase   		=				self.classifier(_A		)

				UpperCAmelCase   		=				None

				if labels is not None:
							if self.config.problem_type is None:
										if self.num_labels == 1:
													UpperCAmelCase   		=				'''regression'''
										elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
													UpperCAmelCase   		=				'''single_label_classification'''
										else:
													UpperCAmelCase   		=				'''multi_label_classification'''
							if self.config.problem_type == "regression":
										UpperCAmelCase   		=				MSELoss()
										if self.num_labels == 1:
													UpperCAmelCase   		=				loss_fct(logits.squeeze()							,					labels.squeeze()		)
										else:
													UpperCAmelCase   		=				loss_fct(_A							,					_A		)
							elif self.config.problem_type == "single_label_classification":
										UpperCAmelCase   		=				CrossEntropyLoss()
										UpperCAmelCase   		=				loss_fct(logits.view(-1							,					self.num_labels		)							,					labels.view(-1		)		)
							elif self.config.problem_type == "multi_label_classification":
										UpperCAmelCase   		=				BCEWithLogitsLoss()
										UpperCAmelCase   		=				loss_fct(_A							,					_A		)

				if not return_dict:
							UpperCAmelCase   		=				(logits,) + outputs[2:]
							return (loss,) + output if loss is not None else output

				return ImageClassifierOutputWithNoAttention(loss=_A							,					logits=_A							,					hidden_states=outputs.hidden_states		)

@add_start_docstrings(
    '''
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    '''     ,							a_     ,							)
class 	A_       (a_     ,							a_			):
	def __init__(    self							,					_A		):






				'''simple docstring'''







				super().__init__(_A		)
				super()._init_backbone(_A		)

				UpperCAmelCase   		=				[config.embedding_size] + config.hidden_sizes
				UpperCAmelCase   		=				ResNetEmbeddings(_A		)
				UpperCAmelCase   		=				ResNetEncoder(_A		)

				# initialize weights and apply final processing
				self.post_init()







	@add_start_docstrings_to_model_forward(_A		)
	@replace_return_docstrings(output_type=_A							,					config_class=_CONFIG_FOR_DOC		)
	def      _lowercase				(    self							,					_A							,					_A = None							,					_A = None		):






				'''simple docstring'''







				UpperCAmelCase   		=				return_dict if return_dict is not None else self.config.use_return_dict
				UpperCAmelCase   		=				(
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)

				UpperCAmelCase   		=				self.embedder(_A		)

				UpperCAmelCase   		=				self.encoder(_A							,					output_hidden_states=_A							,					return_dict=_A		)

				UpperCAmelCase   		=				outputs.hidden_states

				UpperCAmelCase   		=				()
				for idx, stage in enumerate(self.stage_names		):
							if stage in self.out_features:
										feature_maps += (hidden_states[idx],)

				if not return_dict:
							UpperCAmelCase   		=				(feature_maps,)
							if output_hidden_states:
										output += (outputs.hidden_states,)
							return output

				return BackboneOutput(
				    feature_maps=_A							,					hidden_states=outputs.hidden_states if output_hidden_states else None							,					attentions=_A							,					)










import argparse
import re

from flax.traverse_util import flatten_dict, unflatten_dict
from tax import checkpoints

from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration
from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model
from transformers.utils import logging


logging.set_verbosity_info()


# should not include what is already done by the `from_pt` argument
__A				:    int       =   {
    "/attention/": "/0/SelfAttention/",
    "/self_attention/": "/0/SelfAttention/",
    "/encoder_decoder_attention/": "/1/EncDecAttention/",
    "value": "v",
    "query": "q",
    "key": "k",
    "out": "o",
    "pre_self_attention_layer_norm": "0/layer_norm",
    "pre_cross_attention_layer_norm": "1/layer_norm",
    "pre_attention_layer_norm": "0/layer_norm",  # previously 1, but seems wrong
    "token_embedder": "shared",
    "encoder_norm": "final_layer_norm",
    "decoder_norm": "final_layer_norm",
    "relpos_bias/rel_embedding": "block/0/layer/0/SelfAttention/relative_attention_bias/weight",
    "router/router_weights/w/": "router/classifier/",
    "roer/roer_weights/w/": "router/classifier/",
    "logits_dense": "lm_head",
}





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				List[str]:







			'''simple docstring'''

			UpperCAmelCase   		=				list(s_dict.keys() )
			for key in keys:
						UpperCAmelCase   		=				R'''.*/layers_(\d+)'''
						UpperCAmelCase   		=				key
						if re.match(UpperCamelCase__			,			UpperCamelCase__ ):
									UpperCAmelCase   		=				re.sub(R'''layers_(\d+)'''			,			R'''block/\1/layer'''			,			UpperCamelCase__ )

						UpperCAmelCase   		=				R'''(encoder|decoder)\/'''

						if re.match(UpperCamelCase__			,			UpperCamelCase__ ):
									UpperCAmelCase   		=				re.match(UpperCamelCase__			,			UpperCamelCase__ ).groups()
									if groups[0] == "encoder":
												UpperCAmelCase   		=				re.sub(R'''/mlp/'''			,			R'''/1/mlp/'''			,			UpperCamelCase__ )
												UpperCAmelCase   		=				re.sub(R'''/pre_mlp_layer_norm/'''			,			R'''/1/layer_norm/'''			,			UpperCamelCase__ )

									elif groups[0] == "decoder":
												UpperCAmelCase   		=				re.sub(R'''/mlp/'''			,			R'''/2/mlp/'''			,			UpperCamelCase__ )
												UpperCAmelCase   		=				re.sub(R'''/pre_mlp_layer_norm/'''			,			R'''/2/layer_norm/'''			,			UpperCamelCase__ )

        # 2. Convert other classic mappings
						for old_key, temp_key in MOE_LAYER_NAME_MAPPING.items():
									if old_key in new_key:
												UpperCAmelCase   		=				new_key.replace(UpperCamelCase__			,			UpperCamelCase__ )

						print(F"""{key} -> {new_key}""" )
						UpperCAmelCase   		=				s_dict.pop(UpperCamelCase__ )

			if "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
						UpperCAmelCase   		=				s_dict[
						    '''encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight'''
						].T
			if "decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
						UpperCAmelCase   		=				s_dict[
						    '''decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight'''
						].T

			# 3. Take extra care of the EXPERTS layer
			for key in list(s_dict.keys() ):
						if "expert" in key:
									UpperCAmelCase   		=				s_dict[key].shape[0]
									UpperCAmelCase   		=				s_dict[key]
									for idx in range(UpperCamelCase__ ):
												UpperCAmelCase   		=				expert_weihts[idx]
												print(F"""{key} -> {key.replace("expert/"			,			"nested fstring" )}""" )

									s_dict.pop(UpperCamelCase__ )

			return s_dict


__A				:    Optional[int]       =   {
    "NUM_ENCODER_LAYERS": "num_layers",
    "NUM_DECODER_LAYERS": "num_decoder_layers",
    "NUM_HEADS": "num_heads",
    "HEAD_DIM": "d_kv",
    "EMBED_DIM": "d_model",
    "MLP_DIM": "d_ff",
    "NUM_SELECTED_EXPERTS": "num_selected_experts",
    "NUM_ENCODER_SPARSE_LAYERS": "num_sparse_encoder_layers",
    "NUM_DECODER_SPARSE_LAYERS": "num_sparse_decoder_layers",
    "dense.MlpBlock.activations": "feed_forward_proj",
}





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__			,			UpperCamelCase__ )  ->				List[str]:







			'''simple docstring'''

			import regex as re

			with open(UpperCamelCase__			,			'''r''' ) as f:
						UpperCAmelCase   		=				f.read()

			UpperCAmelCase   		=				re.findall(R'''(.*) = ([0-9.]*)'''			,			UpperCamelCase__ )
			UpperCAmelCase   		=				{}
			for param, value in regex_match:
						if param in GIN_TO_CONFIG_MAPPING and value != "":
									UpperCAmelCase   		=				float(UpperCamelCase__ ) if '''.''' in value else int(UpperCamelCase__ )

			UpperCAmelCase   		=				re.findall(R'''(.*activations) = \(\'(.*)\',\)'''			,			UpperCamelCase__ )[0]
			UpperCAmelCase   		=				str(activation[1] )

			UpperCAmelCase   		=				num_experts
			UpperCAmelCase   		=				SwitchTransformersConfig(**UpperCamelCase__ )
			return config





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__=None			,			UpperCamelCase__="./"			,			UpperCamelCase__=8 )  ->				List[Any]:







			'''simple docstring'''

			print(F"""Loading flax weights from : {flax_checkpoint_path}""" )
			UpperCAmelCase   		=				checkpoints.load_tax_checkpoint(UpperCamelCase__ )

			if gin_file is not None:
						UpperCAmelCase   		=				convert_gin_to_config(UpperCamelCase__			,			UpperCamelCase__ )
			else:
						UpperCAmelCase   		=				SwitchTransformersConfig.from_pretrained(UpperCamelCase__ )

			UpperCAmelCase   		=				SwitchTransformersForConditionalGeneration(UpperCamelCase__ )

			UpperCAmelCase   		=				flax_params['''target''']
			UpperCAmelCase   		=				flatten_dict(UpperCamelCase__			,			sep='''/''' )
			UpperCAmelCase   		=				rename_keys(UpperCamelCase__ )
			UpperCAmelCase   		=				unflatten_dict(UpperCamelCase__			,			sep='''/''' )

			# Load the flax params in the PT model
			load_flax_weights_in_pytorch_model(UpperCamelCase__			,			UpperCamelCase__ )

			print(F"""Save PyTorch model to {pytorch_dump_path}""" )
			pt_model.save_pretrained(UpperCamelCase__ )


if __name__ == "__main__":
				__A				:    Tuple       =   argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    "--switch_t5x_checkpoint_path",
				    default=None,
				    type=str,
				    required=True,
				    help=(
				        "The config json file corresponding to the pre-trained SwitchTransformers model. \nThis specifies the"
				        " model architecture. If not provided, a `gin_file` has to be provided."
				    ),
				)
				parser.add_argument(
				    "--gin_file",
				    default=None,
				    type=str,
				    required=False,
				    help="Path to the gin config file. If not provided, a `config_file` has to be passed   ",
				)
				parser.add_argument(
				    "--config_name", default=None, type=str, required=False, help="Config name of SwitchTransformers model."
				)
				parser.add_argument(
				    "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output pytorch model."
				)
				parser.add_argument("--num_experts", default=8, type=int, required=False, help="Number of experts")
				__A				:    Tuple       =   parser.parse_args()
				convert_flax_checkpoint_to_pytorch(
				    args.switch_tax_checkpoint_path,
				    args.config_name,
				    args.gin_file,
				    args.pytorch_dump_folder_path,
				    args.num_experts,
				)



