





'''simple docstring'''







from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_a					: List[Any]      =    logging.get_logger(__name__)

_a					: str      =    {
    """mnaylor/mega-base-wikitext""": """https://huggingface.co/mnaylor/mega-base-wikitext/resolve/main/config.json""",
}





class     _UpperCAmelCase							(						lowerCAmelCase_   ):
		a   :   List[Any]     ="""mega"""
		def __init__(  self,__SCREAMING_SNAKE_CASE=3_05_22,__SCREAMING_SNAKE_CASE=1_28,__SCREAMING_SNAKE_CASE=4,__SCREAMING_SNAKE_CASE=2_56,__SCREAMING_SNAKE_CASE=16,__SCREAMING_SNAKE_CASE=True,__SCREAMING_SNAKE_CASE=64,__SCREAMING_SNAKE_CASE=False,__SCREAMING_SNAKE_CASE=-1,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=True,__SCREAMING_SNAKE_CASE="scalenorm",__SCREAMING_SNAKE_CASE=True,__SCREAMING_SNAKE_CASE="silu",__SCREAMING_SNAKE_CASE="softmax",__SCREAMING_SNAKE_CASE=0.1,__SCREAMING_SNAKE_CASE=0.1,__SCREAMING_SNAKE_CASE=0.1,__SCREAMING_SNAKE_CASE=False,__SCREAMING_SNAKE_CASE=True,__SCREAMING_SNAKE_CASE=2_56,__SCREAMING_SNAKE_CASE=True,__SCREAMING_SNAKE_CASE=0.1,__SCREAMING_SNAKE_CASE=20_48,__SCREAMING_SNAKE_CASE=False,__SCREAMING_SNAKE_CASE=2,__SCREAMING_SNAKE_CASE=0.02,__SCREAMING_SNAKE_CASE=0.2,__SCREAMING_SNAKE_CASE=0.02,__SCREAMING_SNAKE_CASE=1.0,__SCREAMING_SNAKE_CASE=1,__SCREAMING_SNAKE_CASE=0,__SCREAMING_SNAKE_CASE=2,__SCREAMING_SNAKE_CASE="rotary",__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=True,__SCREAMING_SNAKE_CASE=True,**__SCREAMING_SNAKE_CASE,):



					'''simple docstring'''

					super().__init__(pad_token_id=__SCREAMING_SNAKE_CASE,bos_token_id=__SCREAMING_SNAKE_CASE,eos_token_id=__SCREAMING_SNAKE_CASE,**__SCREAMING_SNAKE_CASE			)

					__lowerCAmelCase											=				vocab_size
					__lowerCAmelCase											=				hidden_size
					__lowerCAmelCase											=				num_hidden_layers
					__lowerCAmelCase											=				activation
					__lowerCAmelCase											=				attention_activation
					__lowerCAmelCase											=				intermediate_size
					__lowerCAmelCase											=				ema_projection_size
					__lowerCAmelCase											=				bidirectional
					__lowerCAmelCase											=				shared_representation_size
					__lowerCAmelCase											=				use_chunking
					__lowerCAmelCase											=				chunk_size
					__lowerCAmelCase											=				truncation
					__lowerCAmelCase											=				normalize_before_mega
					__lowerCAmelCase											=				normalization_type
					__lowerCAmelCase											=				norm_affine
					__lowerCAmelCase											=				dropout_prob
					__lowerCAmelCase											=				hidden_dropout_prob
					__lowerCAmelCase											=				attention_probs_dropout_prob
					__lowerCAmelCase											=				use_feature_dropout
					__lowerCAmelCase											=				use_normalized_ffn
					__lowerCAmelCase											=				nffn_hidden_size
					__lowerCAmelCase											=				normalize_before_ffn
					__lowerCAmelCase											=				nffn_activation_dropout_prob
					__lowerCAmelCase											=				max_positions
					__lowerCAmelCase											=				add_token_type_embeddings
					__lowerCAmelCase											=				type_vocab_size
					__lowerCAmelCase											=				initializer_range
					__lowerCAmelCase											=				ema_delta_alpha_range
					__lowerCAmelCase											=				ema_beta_range
					__lowerCAmelCase											=				ema_gamma_omega_range
					__lowerCAmelCase											=				relative_positional_bias
					__lowerCAmelCase											=				use_cache
					__lowerCAmelCase											=				classifier_dropout
					__lowerCAmelCase											=				add_lm_hidden_dense_layer
					__lowerCAmelCase											=				1  # not used but required by Hugging Face





class     _UpperCAmelCase							(						lowerCAmelCase_   ):
		@property
		def   lowerCamelCase__			(  self			):



					'''simple docstring'''

					if self.task == "multiple-choice":
								__lowerCAmelCase											=				{0: """batch""", 1: """choice""", 2: """sequence"""}
					else:
								__lowerCAmelCase											=				{0: """batch""", 1: """sequence"""}
					return OrderedDict(
					    [
					        ("""input_ids""", dynamic_axis),
					        ("""attention_mask""", dynamic_axis),
					    ]			)








'''simple docstring'''







import torch
from torch import autograd





class     _UpperCAmelCase							(						autograd.Function   ):
		@staticmethod
		def   lowerCamelCase__			(  __SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE			):



					'''simple docstring'''

					__lowerCAmelCase											=				inputs.numel()
					__lowerCAmelCase											=				int(0.005 * nb_elems			) + 1
					if sigmoid:
								__lowerCAmelCase											=				(torch.sigmoid(__SCREAMING_SNAKE_CASE			) > threshold).type(inputs.type()			)
					else:
								__lowerCAmelCase											=				(inputs > threshold).type(inputs.type()			)
					if mask.sum() < nb_min:
								# We limit the pruning so that at least 0.5% (half a percent) of the weights are remaining
								__lowerCAmelCase											=				inputs.flatten().kthvalue(max(nb_elems - nb_min,1			)			).values
								__lowerCAmelCase											=				(inputs > k_threshold).type(inputs.type()			)
					return mask







		@staticmethod
		def   lowerCamelCase__			(  __SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE			):



					'''simple docstring'''

					return gradOutput, None, None





class     _UpperCAmelCase							(						autograd.Function   ):
		@staticmethod
		def   lowerCamelCase__			(  __SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE			):



					'''simple docstring'''

					__lowerCAmelCase											=				inputs.clone()
					__lowerCAmelCase						,       __lowerCAmelCase											=				inputs.flatten().sort(descending=__SCREAMING_SNAKE_CASE			)
					__lowerCAmelCase											=				int(threshold * inputs.numel()			)

					# flat_out and mask access the same memory.
					__lowerCAmelCase											=				mask.flatten()
					__lowerCAmelCase											=				0
					__lowerCAmelCase											=				1
					return mask







		@staticmethod
		def   lowerCamelCase__			(  __SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE			):



					'''simple docstring'''

					return gradOutput, None





class     _UpperCAmelCase							(						lowerCAmelCase_   ):
		@staticmethod
		def   lowerCamelCase__			(  __SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE			):



					'''simple docstring'''

					__lowerCAmelCase											=				inputs.clone()
					__lowerCAmelCase						,       __lowerCAmelCase											=				inputs.abs().flatten().sort(descending=__SCREAMING_SNAKE_CASE			)
					__lowerCAmelCase											=				int(threshold * inputs.numel()			)

					# flat_out and mask access the same memory.
					__lowerCAmelCase											=				mask.flatten()
					__lowerCAmelCase											=				0
					__lowerCAmelCase											=				1
					return mask

