




'''simple docstring'''




import torch
from torch import autograd






class      UpperCamelCase_  (autograd.Function  ):







						"""simple docstring"""


						@staticmethod
						def    _a     (  _lowerCamelCase   :							Optional[int]				,			_lowerCamelCase   :							torch.tensor				,			_lowerCamelCase   :							float				,			_lowerCamelCase   :							bool     ):





											"""simple docstring"""







											A_			:  int			=				inputs.numel()
											A_			:  str			=				int(0.0_05 * nb_elems     ) + 1
											if sigmoid:
																A_			:  List[str]			=				(torch.sigmoid(_lowerCamelCase     ) > threshold).type(inputs.type()     )
											else:
																A_			:  Optional[Any]			=				(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
																A_			:  Union[str, Any]			=				inputs.flatten().kthvalue(max(nb_elems - nb_min				,			1     )     ).values
																A_			:  List[str]			=				(inputs > k_threshold).type(inputs.type()     )
											return mask







						@staticmethod
						def    _a     (  _lowerCamelCase   :							Union[str, Any]				,			_lowerCamelCase   :							int     ):





											"""simple docstring"""







											return gradOutput, None, None






class      UpperCamelCase_  (autograd.Function  ):







						"""simple docstring"""


						@staticmethod
						def    _a     (  _lowerCamelCase   :							Tuple				,			_lowerCamelCase   :							torch.tensor				,			_lowerCamelCase   :							float     ):





											"""simple docstring"""







											A_			:  Optional[int]			=				inputs.clone()
											A_			,A_			:  str			=				inputs.flatten().sort(descending=_lowerCamelCase     )
											A_			:  Optional[int]			=				int(threshold * inputs.numel()     )

											# flat_out and mask access the same memory.
											A_			:  Optional[int]			=				mask.flatten()
											A_			:  Tuple			=				0
											A_			:  Optional[Any]			=				1
											return mask







						@staticmethod
						def    _a     (  _lowerCamelCase   :							Union[str, Any]				,			_lowerCamelCase   :							str     ):





											"""simple docstring"""







											return gradOutput, None






class      UpperCamelCase_  (a__  ):







						"""simple docstring"""


						@staticmethod
						def    _a     (  _lowerCamelCase   :							torch.tensor				,			_lowerCamelCase   :							float     ):





											"""simple docstring"""







											A_			:  Any			=				inputs.clone()
											A_			,A_			:  Dict			=				inputs.abs().flatten().sort(descending=_lowerCamelCase     )
											A_			:  Dict			=				int(threshold * inputs.numel()     )

											# flat_out and mask access the same memory.
											A_			:  str			=				mask.flatten()
											A_			:  int			=				0
											A_			:  str			=				1
											return mask









'''simple docstring'''




from ...configuration_utils import PretrainedConfig
from ...utils import logging


snake_case__ 							=    logging.get_logger(__name__)

snake_case__ 							=    {
    """uw-madison/nystromformer-512""": """https://huggingface.co/uw-madison/nystromformer-512/resolve/main/config.json""",
    # See all Nystromformer models at https://huggingface.co/models?filter=nystromformer
}






class      UpperCamelCase_  (a__  ):







						"""simple docstring"""


						_lowerCAmelCase         =       'nystromformer'







						def __init__(  self   :							Optional[Any]				,			_lowerCamelCase   :							str=30000				,			_lowerCamelCase   :							Optional[Any]=768				,			_lowerCamelCase   :							Optional[Any]=12				,			_lowerCamelCase   :							str=12				,			_lowerCamelCase   :							List[str]=3072				,			_lowerCamelCase   :							Any="gelu_new"				,			_lowerCamelCase   :							List[Any]=0.1				,			_lowerCamelCase   :							List[str]=0.1				,			_lowerCamelCase   :							List[str]=510				,			_lowerCamelCase   :							Optional[Any]=2				,			_lowerCamelCase   :							List[Any]=64				,			_lowerCamelCase   :							str=64				,			_lowerCamelCase   :							str=65				,			_lowerCamelCase   :							List[Any]=False				,			_lowerCamelCase   :							List[Any]=0.02				,			_lowerCamelCase   :							List[Any]=1E-5				,			_lowerCamelCase   :							Dict=1				,			_lowerCamelCase   :							Tuple=0				,			_lowerCamelCase   :							List[str]=2				,			**_lowerCamelCase   :							List[str]				,			):





											"""simple docstring"""







											A_			:  List[Any]			=				vocab_size
											A_			:  Tuple			=				max_position_embeddings
											A_			:  Dict			=				hidden_size
											A_			:  Optional[Any]			=				num_hidden_layers
											A_			:  Dict			=				num_attention_heads
											A_			:  Tuple			=				intermediate_size
											A_			:  Dict			=				hidden_act
											A_			:  List[Any]			=				hidden_dropout_prob
											A_			:  str			=				attention_probs_dropout_prob
											A_			:  Optional[Any]			=				initializer_range
											A_			:  List[str]			=				type_vocab_size
											A_			:  List[Any]			=				segment_means_seq_len
											A_			:  Dict			=				num_landmarks
											A_			:  int			=				conv_kernel_size
											A_			:  Optional[Any]			=				inv_coeff_init_option
											A_			:  str			=				layer_norm_eps
											super().__init__(pad_token_id=_lowerCamelCase				,			bos_token_id=_lowerCamelCase				,			eos_token_id=_lowerCamelCase				,			**_lowerCamelCase     )



