


'''simple docstring'''






import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


__UpperCAmelCase    				=							[nn.LayerNorm]

__UpperCAmelCase    				=							logging.get_logger(__name__)

__UpperCAmelCase    				=							version.parse(version.parse(torch.__version__).base_version)

__UpperCAmelCase    				=							parsed_torch_version_base >= version.parse('''2.0''')
__UpperCAmelCase    				=							parsed_torch_version_base >= version.parse('''1.12''')
__UpperCAmelCase    				=							parsed_torch_version_base >= version.parse('''1.11''')
__UpperCAmelCase    				=							parsed_torch_version_base < version.parse('''1.11''')


def 					_snake_case   (     A ,    A ,    A ,    A ,    self       )						->    int:
							from torch import _softmax_backward_data

							if is_torch_less_than_1_11:
														return _softmax_backward_data(A ,    A ,    parent.dim ,    self       )
							else:
														return _softmax_backward_data(A ,    A ,    parent.dim ,    self.dtype       )


def 					_snake_case   (     A ,    A ,    A = 0       )						->    nn.Linear:
							lowerCAmelCase__          =	index.to(layer.weight.device       )
							lowerCAmelCase__          =	layer.weight.index_select(A ,    A       ).clone().detach()
							if layer.bias is not None:
														if dim == 1:
																					lowerCAmelCase__          =	layer.bias.clone().detach()
														else:
																					lowerCAmelCase__          =	layer.bias[index].clone().detach()
							lowerCAmelCase__          =	list(layer.weight.size()       )
							lowerCAmelCase__          =	len(A       )
							lowerCAmelCase__          =	nn.Linear(new_size[1] ,    new_size[0] ,    bias=layer.bias is not None       ).to(layer.weight.device       )
							lowerCAmelCase__          =	False
							new_layer.weight.copy_(W.contiguous()       )
							lowerCAmelCase__          =	True
							if layer.bias is not None:
														lowerCAmelCase__          =	False
														new_layer.bias.copy_(b.contiguous()       )
														lowerCAmelCase__          =	True
							return new_layer




class 		a__				(							nn.Module  ):





							'''simple docstring'''






							def __init__(		self					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							Tuple:
														super().__init__()
														lowerCAmelCase__          =	nf
														lowerCAmelCase__          =	nn.Parameter(torch.empty(lowerCamelCase_					,    lowerCamelCase_   )   )
														lowerCAmelCase__          =	nn.Parameter(torch.zeros(lowerCamelCase_   )   )
														nn.init.normal_(self.weight					,    std=0.02   )







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_   )       ->							List[str]:
														lowerCAmelCase__          =	x.size()[:-1] + (self.nf,)
														lowerCAmelCase__          =	torch.addmm(self.bias					,    x.view(-1					,    x.size(-1   )   )					,    self.weight   )
														lowerCAmelCase__          =	x.view(lowerCamelCase_   )
														return x


def 					_snake_case   (     A ,    A ,    A = 1       )						->    ConvaD:
							lowerCAmelCase__          =	index.to(layer.weight.device       )
							lowerCAmelCase__          =	layer.weight.index_select(A ,    A       ).clone().detach()
							if dim == 0:
														lowerCAmelCase__          =	layer.bias.clone().detach()
							else:
														lowerCAmelCase__          =	layer.bias[index].clone().detach()
							lowerCAmelCase__          =	list(layer.weight.size()       )
							lowerCAmelCase__          =	len(A       )
							lowerCAmelCase__          =	ConvaD(new_size[1] ,    new_size[0]       ).to(layer.weight.device       )
							lowerCAmelCase__          =	False
							new_layer.weight.copy_(W.contiguous()       )
							lowerCAmelCase__          =	True
							lowerCAmelCase__          =	False
							new_layer.bias.copy_(b.contiguous()       )
							lowerCAmelCase__          =	True
							return new_layer


def 					_snake_case   (     A ,    A ,    A = None       )						->    Union[nn.Linear, ConvaD]:
							if isinstance(A ,    nn.Linear       ):
														return prune_linear_layer(A ,    A ,    dim=0 if dim is None else dim       )
							elif isinstance(A ,    A       ):
														return prune_convad_layer(A ,    A ,    dim=1 if dim is None else dim       )
							else:
														raise ValueError(F"""Can't prune layer of class {layer.__class__}"""       )


def 					_snake_case   (     A ,    A ,    A ,    *A       )						->    torch.Tensor:
							assert len(A       ) > 0, F"""{input_tensors} has to be a tuple/list of tensors"""

							# inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
							lowerCAmelCase__          =	len(inspect.signature(A       ).parameters       )
							if num_args_in_forward_chunk_fn != len(A       ):
														raise ValueError(
														    F"""forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(A       )} input """
														    '''tensors are given'''       )

							if chunk_size > 0:
														lowerCAmelCase__          =	input_tensors[0].shape[chunk_dim]
														for input_tensor in input_tensors:
																					if input_tensor.shape[chunk_dim] != tensor_shape:
																												raise ValueError(
																												    F"""All input tenors have to be of the same shape: {tensor_shape}, """
																												    F"""found shape {input_tensor.shape[chunk_dim]}"""       )

														if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
																					raise ValueError(
																					    F"""The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk """
																					    F"""size {chunk_size}"""       )

														lowerCAmelCase__          =	input_tensors[0].shape[chunk_dim] // chunk_size

														# chunk input tensor into tuples
														lowerCAmelCase__          =	tuple(input_tensor.chunk(A ,    dim=A       ) for input_tensor in input_tensors       )
														# apply forward fn to every tuple
														lowerCAmelCase__          =	tuple(forward_fn(*A       ) for input_tensors_chunk in zip(*A       )       )
														# concatenate output at same dimension
														return torch.cat(A ,    dim=A       )

							return forward_fn(*A       )


def 					_snake_case   (     A ,    A ,    A ,    A       )						->    Tuple[Set[int], torch.LongTensor]:
							lowerCAmelCase__          =	torch.ones(A ,    A       )
							lowerCAmelCase__          =	set(A       ) - already_pruned_heads  # Convert to set and remove already pruned heads
							for head in heads:
														# Compute how many pruned heads are before the head and move the index accordingly
														lowerCAmelCase__          =	head - sum(1 if h < head else 0 for h in already_pruned_heads       )
														lowerCAmelCase__          =	0
							lowerCAmelCase__          =	mask.view(-1       ).contiguous().eq(1       )
							lowerCAmelCase__          =	torch.arange(len(A       )       )[mask].long()
							return heads, index


def 					_snake_case   (     *A ,    A = None       )						->    Tuple[torch.Tensor, ...]:
							return torch.meshgrid(*A ,    indexing=A       )






def 					_snake_case   (     A       )						->    Tuple[torch.device, int, int]:
							return tensor.device, storage_ptr(A       ), storage_size(A       )



'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase    				=							{
    '''configuration_switch_transformers''': [
        '''SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''SwitchTransformersConfig''',
        '''SwitchTransformersOnnxConfig''',
    ]
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__UpperCAmelCase    				=							[
		    '''SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''SwitchTransformersEncoderModel''',
		    '''SwitchTransformersForConditionalGeneration''',
		    '''SwitchTransformersModel''',
		    '''SwitchTransformersPreTrainedModel''',
		    '''SwitchTransformersTop1Router''',
		    '''SwitchTransformersSparseMLP''',
		]


if TYPE_CHECKING:
		from .configuration_switch_transformers import (
		    SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    SwitchTransformersConfig,
		    SwitchTransformersOnnxConfig,
		)

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_switch_transformers import (
				    SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST,
				    SwitchTransformersEncoderModel,
				    SwitchTransformersForConditionalGeneration,
				    SwitchTransformersModel,
				    SwitchTransformersPreTrainedModel,
				    SwitchTransformersSparseMLP,
				    SwitchTransformersTopaRouter,
				)


else:
		import sys

		__UpperCAmelCase    				=							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)