
'''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


_lowercase					:    List[str]       				=			[nn.LayerNorm]

_lowercase					:    Tuple       				=			logging.get_logger(__name__)

_lowercase					:    int       				=			version.parse(version.parse(torch.__version__).base_version)

_lowercase					:    Any       				=			parsed_torch_version_base >= version.parse("2.0")
_lowercase					:    List[str]       				=			parsed_torch_version_base >= version.parse("1.12")
_lowercase					:    List[str]       				=			parsed_torch_version_base >= version.parse("1.11")
_lowercase					:    Union[str, Any]       				=			parsed_torch_version_base < version.parse("1.11")
def    snake_case_						(							__SCREAMING_SNAKE_CASE   :					Optional[int] , __SCREAMING_SNAKE_CASE   :					str , __SCREAMING_SNAKE_CASE   :					List[str] , __SCREAMING_SNAKE_CASE   :					Optional[Any] , self   :					Optional[int]				):



						"""simple docstring"""





						from torch import _softmax_backward_data

						if is_torch_less_than_1_11:
												return _softmax_backward_data(lowerCAmelCase_ , lowerCAmelCase_ , parent.dim , self				)
						else:
												return _softmax_backward_data(lowerCAmelCase_ , lowerCAmelCase_ , parent.dim , self.dtype				)
def    snake_case_						(							__SCREAMING_SNAKE_CASE   :					str , __SCREAMING_SNAKE_CASE   :					Dict , __SCREAMING_SNAKE_CASE   :					Optional[Any] = 0				):



						"""simple docstring"""





						lowercase_      :	List[str]						=			index.to(layer.weight.device				)
						lowercase_      :	Optional[int]						=			layer.weight.index_select(lowerCAmelCase_ , lowerCAmelCase_				).clone().detach()
						if layer.bias is not None:
												if dim == 1:
																		lowercase_      :	Optional[int]						=			layer.bias.clone().detach()
												else:
																		lowercase_      :	Dict						=			layer.bias[index].clone().detach()
						lowercase_      :	List[str]						=			list(layer.weight.size()				)
						lowercase_      :	Any						=			len(lowerCAmelCase_				)
						lowercase_      :	List[Any]						=			nn.Linear(new_size[1] , new_size[0] , bias=layer.bias is not None				).to(layer.weight.device				)
						lowercase_      :	Union[str, Any]						=			False
						new_layer.weight.copy_(W.contiguous()				)
						lowercase_      :	str						=			True
						if layer.bias is not None:
												lowercase_      :	Union[str, Any]						=			False
												new_layer.bias.copy_(b.contiguous()				)
												lowercase_      :	Any						=			True
						return new_layer






class 					lowerCAmelCase__  (    nn.Module	):






					def __init__(  self							,				__SCREAMING_SNAKE_CASE							,				__SCREAMING_SNAKE_CASE						):


											"""simple docstring"""





											super().__init__()
											lowercase_      :	Optional[int]						=			nf
											lowercase_      :	str						=			nn.Parameter(torch.empty(__UpperCamelCase							,				__UpperCamelCase						)						)
											lowercase_      :	Optional[int]						=			nn.Parameter(torch.zeros(__UpperCamelCase						)						)
											nn.init.normal_(self.weight							,				std=0.02						)






					def       _snake_case       (  self							,				__SCREAMING_SNAKE_CASE						):


											"""simple docstring"""





											lowercase_      :	str						=			x.size()[:-1] + (self.nf,)
											lowercase_      :	Any						=			torch.addmm(self.bias							,				x.view(-1							,				x.size(-1						)						)							,				self.weight						)
											lowercase_      :	str						=			x.view(__UpperCamelCase						)
											return x







def    snake_case_						(							__SCREAMING_SNAKE_CASE   :					Optional[int] , __SCREAMING_SNAKE_CASE   :					Dict , __SCREAMING_SNAKE_CASE   :					Optional[int] = 1				):



						"""simple docstring"""





						lowercase_      :	str						=			index.to(layer.weight.device				)
						lowercase_      :	Tuple						=			layer.weight.index_select(lowerCAmelCase_ , lowerCAmelCase_				).clone().detach()
						if dim == 0:
												lowercase_      :	Tuple						=			layer.bias.clone().detach()
						else:
												lowercase_      :	Tuple						=			layer.bias[index].clone().detach()
						lowercase_      :	Optional[int]						=			list(layer.weight.size()				)
						lowercase_      :	List[str]						=			len(lowerCAmelCase_				)
						lowercase_      :	Dict						=			ConvaD(new_size[1] , new_size[0]				).to(layer.weight.device				)
						lowercase_      :	Dict						=			False
						new_layer.weight.copy_(W.contiguous()				)
						lowercase_      :	int						=			True
						lowercase_      :	Optional[int]						=			False
						new_layer.bias.copy_(b.contiguous()				)
						lowercase_      :	Tuple						=			True
						return new_layer
def    snake_case_						(							__SCREAMING_SNAKE_CASE   :					Tuple , __SCREAMING_SNAKE_CASE   :					Optional[int] , __SCREAMING_SNAKE_CASE   :					Dict = None				):



						"""simple docstring"""





						if isinstance(lowerCAmelCase_ , nn.Linear				):
												return prune_linear_layer(lowerCAmelCase_ , lowerCAmelCase_ , dim=0 if dim is None else dim				)
						elif isinstance(lowerCAmelCase_ , lowerCAmelCase_				):
												return prune_convad_layer(lowerCAmelCase_ , lowerCAmelCase_ , dim=1 if dim is None else dim				)
						else:
												raise ValueError(F'''Can\'t prune layer of class {layer.__class__}'''				)
def    snake_case_						(							__SCREAMING_SNAKE_CASE   :					List[Any] , __SCREAMING_SNAKE_CASE   :					Dict , __SCREAMING_SNAKE_CASE   :					Optional[Any] , *__SCREAMING_SNAKE_CASE   :					Any				):



						"""simple docstring"""





						assert len(lowerCAmelCase_				) > 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
						lowercase_      :	Dict						=			len(inspect.signature(lowerCAmelCase_				).parameters				)
						if num_args_in_forward_chunk_fn != len(lowerCAmelCase_				):
												raise ValueError(
												    F'''forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(lowerCAmelCase_				)} input '''
												    '''tensors are given'''				)

						if chunk_size > 0:
												lowercase_      :	str						=			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}'''				)

												lowercase_      :	List[str]						=			input_tensors[0].shape[chunk_dim] // chunk_size

												# chunk input tensor into tuples
												lowercase_      :	Optional[int]						=			tuple(input_tensor.chunk(lowerCAmelCase_ , dim=lowerCAmelCase_				) for input_tensor in input_tensors				)
												# apply forward fn to every tuple
												lowercase_      :	Tuple						=			tuple(forward_fn(*lowerCAmelCase_				) for input_tensors_chunk in zip(*lowerCAmelCase_				)				)
												# concatenate output at same dimension
												return torch.cat(lowerCAmelCase_ , dim=lowerCAmelCase_				)

						return forward_fn(*lowerCAmelCase_				)
def    snake_case_						(							__SCREAMING_SNAKE_CASE   :					str , __SCREAMING_SNAKE_CASE   :					List[Any] , __SCREAMING_SNAKE_CASE   :					List[Any] , __SCREAMING_SNAKE_CASE   :					Dict				):



						"""simple docstring"""





						lowercase_      :	Any						=			torch.ones(lowerCAmelCase_ , lowerCAmelCase_				)
						lowercase_      :	Optional[int]						=			set(lowerCAmelCase_				) - 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
												lowercase_      :	int						=			head - sum(1 if h < head else 0 for h in already_pruned_heads				)
												lowercase_      :	Optional[Any]						=			0
						lowercase_      :	str						=			mask.view(-1				).contiguous().eq(1				)
						lowercase_      :	Optional[Any]						=			torch.arange(len(lowerCAmelCase_				)				)[mask].long()
						return heads, index
def    snake_case_						(							*__SCREAMING_SNAKE_CASE   :					List[Any] , __SCREAMING_SNAKE_CASE   :					Optional[Any] = None				):



						"""simple docstring"""





						return torch.meshgrid(*lowerCAmelCase_ , indexing=lowerCAmelCase_				)




def    snake_case_						(							__SCREAMING_SNAKE_CASE   :					Optional[Any]				):



						"""simple docstring"""





						return tensor.device, storage_ptr(lowerCAmelCase_				), storage_size(lowerCAmelCase_				)








'''simple docstring'''







import unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
      from transformers.models.electra.modeling_flax_electra import (
          FlaxElectraForCausalLM,
          FlaxElectraForMaskedLM,
          FlaxElectraForMultipleChoice,
          FlaxElectraForPreTraining,
          FlaxElectraForQuestionAnswering,
          FlaxElectraForSequenceClassification,
          FlaxElectraForTokenClassification,
          FlaxElectraModel,
      )

class    lowerCamelCase_				(						unittest.TestCase	):






   '''simple docstring'''

   def __init__(	self  ,    __UpperCamelCase  ,    __UpperCamelCase=13  ,    __UpperCamelCase=7  ,    __UpperCamelCase=True  ,    __UpperCamelCase=True  ,    __UpperCamelCase=True  ,    __UpperCamelCase=True  ,    __UpperCamelCase=99  ,    __UpperCamelCase=24  ,    __UpperCamelCase=32  ,    __UpperCamelCase=5  ,    __UpperCamelCase=4  ,    __UpperCamelCase=37  ,    __UpperCamelCase="gelu"  ,    __UpperCamelCase=0.1  ,    __UpperCamelCase=0.1  ,    __UpperCamelCase=5_12  ,    __UpperCamelCase=16  ,    __UpperCamelCase=2  ,    __UpperCamelCase=0.02  ,    __UpperCamelCase=4  ,    ):
       UpperCamelCase__	     =      parent
       UpperCamelCase__	     =      batch_size
       UpperCamelCase__	     =      seq_length
       UpperCamelCase__	     =      is_training
       UpperCamelCase__	     =      use_attention_mask
       UpperCamelCase__	     =      use_token_type_ids
       UpperCamelCase__	     =      use_labels
       UpperCamelCase__	     =      vocab_size
       UpperCamelCase__	     =      hidden_size
       UpperCamelCase__	     =      embedding_size
       UpperCamelCase__	     =      num_hidden_layers
       UpperCamelCase__	     =      num_attention_heads
       UpperCamelCase__	     =      intermediate_size
       UpperCamelCase__	     =      hidden_act
       UpperCamelCase__	     =      hidden_dropout_prob
       UpperCamelCase__	     =      attention_probs_dropout_prob
       UpperCamelCase__	     =      max_position_embeddings
       UpperCamelCase__	     =      type_vocab_size
       UpperCamelCase__	     =      type_sequence_label_size
       UpperCamelCase__	     =      initializer_range
       UpperCamelCase__	     =      num_choices


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      ids_tensor([self.batch_size, self.seq_length]  ,    self.vocab_size			)

       UpperCamelCase__	     =      None
       if self.use_attention_mask:
           UpperCamelCase__	     =      random_attention_mask([self.batch_size, self.seq_length]			)

       UpperCamelCase__	     =      None
       if self.use_token_type_ids:
           UpperCamelCase__	     =      ids_tensor([self.batch_size, self.seq_length]  ,    self.type_vocab_size			)

       UpperCamelCase__	     =      ElectraConfig(
           vocab_size=self.vocab_size  ,    hidden_size=self.hidden_size  ,    embedding_size=self.embedding_size  ,    num_hidden_layers=self.num_hidden_layers  ,    num_attention_heads=self.num_attention_heads  ,    intermediate_size=self.intermediate_size  ,    hidden_act=self.hidden_act  ,    hidden_dropout_prob=self.hidden_dropout_prob  ,    attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,    max_position_embeddings=self.max_position_embeddings  ,    type_vocab_size=self.type_vocab_size  ,    initializer_range=self.initializer_range  ,    )

       return config, input_ids, token_type_ids, attention_mask







   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      self.prepare_config_and_inputs()
       UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__     ,	UpperCamelCase__	     =      config_and_inputs
       UpperCamelCase__	     =      {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
       return config, inputs_dict

@require_flax
class    lowerCamelCase_				(						lowerCAmelCase_       ,     unittest.TestCase	):






   '''simple docstring'''

   A			:  Dict					=						True

   A			:  Optional[Any]					=						(
       (
           FlaxElectraModel,
           FlaxElectraForCausalLM,
           FlaxElectraForMaskedLM,
           FlaxElectraForPreTraining,
           FlaxElectraForTokenClassification,
           FlaxElectraForQuestionAnswering,
           FlaxElectraForMultipleChoice,
           FlaxElectraForSequenceClassification,
       )
       if is_flax_available()
       else ()
   )


   def 						_UpperCamelCase							(	self			):
       UpperCamelCase__	     =      FlaxElectraModelTester(self			)







   @slow
   def 						_UpperCamelCase							(	self			):
       for model_class_name in self.all_model_classes:
           if model_class_name == FlaxElectraForMaskedLM:
               UpperCamelCase__	     =      model_class_name.from_pretrained("google/electra-small-generator"			)
           else:
               UpperCamelCase__	     =      model_class_name.from_pretrained("google/electra-small-discriminator"			)
           UpperCamelCase__	     =      model(np.ones((1, 1)			)			)
           self.assertIsNotNone(__UpperCamelCase			)




