




'''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_		:    Union[str, Any]									=		[nn.LayerNorm]

UpperCamelCase_		:    Union[str, Any]									=		logging.get_logger(__name__)

UpperCamelCase_		:    Optional[Any]									=		version.parse(version.parse(torch.__version__).base_version)

UpperCamelCase_		:    Optional[Any]									=		parsed_torch_version_base >= version.parse('''2.0''')
UpperCamelCase_		:    str									=		parsed_torch_version_base >= version.parse('''1.12''')
UpperCamelCase_		:    int									=		parsed_torch_version_base >= version.parse('''1.11''')
UpperCamelCase_		:    int									=		parsed_torch_version_base < version.parse('''1.11''')





def 			__a			(     _UpperCamelCase:					Optional[Any]      ,		_UpperCamelCase:					str      ,		_UpperCamelCase:					Tuple      ,		_UpperCamelCase:					Dict      ,		self:					str					)		->				List[Any]:

     """simple docstring"""




     from torch import _softmax_backward_data

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





def 			__a			(     _UpperCamelCase:					nn.Linear      ,		_UpperCamelCase:					torch.LongTensor      ,		_UpperCamelCase:					int = 0					)		->				nn.Linear:

     """simple docstring"""




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

class      _a       (		nn.Module							):






     def __init__(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)							->    List[Any]:
          super().__init__()
          _snake_case							=   nf
          _snake_case							=   nn.Parameter(torch.empty(_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)					)
          _snake_case							=   nn.Parameter(torch.zeros(_SCREAMING_SNAKE_CASE					)					)
          nn.init.normal_(self.weight      ,std=0.0_2					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    int:
          _snake_case							=   x.size()[:-1] + (self.nf,)
          _snake_case							=   torch.addmm(self.bias      ,x.view(-1      ,x.size(-1					)					)      ,self.weight					)
          _snake_case							=   x.view(_SCREAMING_SNAKE_CASE					)
          return x







def 			__a			(     _UpperCamelCase:					ConvaD      ,		_UpperCamelCase:					torch.LongTensor      ,		_UpperCamelCase:					int = 1					)		->				ConvaD:

     """simple docstring"""




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





def 			__a			(     _UpperCamelCase:					Union[nn.Linear, ConvaD]      ,		_UpperCamelCase:					torch.LongTensor      ,		_UpperCamelCase:					Optional[int] = None					)		->				Union[nn.Linear, ConvaD]:

     """simple docstring"""




     if isinstance(_UpperCamelCase      ,		nn.Linear					):
          return prune_linear_layer(_UpperCamelCase      ,		_UpperCamelCase      ,		dim=0 if dim is None else dim					)
     elif isinstance(_UpperCamelCase      ,		_UpperCamelCase					):
          return prune_convad_layer(_UpperCamelCase      ,		_UpperCamelCase      ,		dim=1 if dim is None else dim					)
     else:
          raise ValueError(F"""Can't prune layer of class {layer.__class__}"""					)





def 			__a			(     _UpperCamelCase:					Callable[..., torch.Tensor]      ,		_UpperCamelCase:					int      ,		_UpperCamelCase:					int      ,		*_UpperCamelCase:					int					)		->				torch.Tensor:

     """simple docstring"""




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

     if chunk_size > 0:
          _snake_case							=   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}"""					)

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

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

     return forward_fn(*_UpperCamelCase					)





def 			__a			(     _UpperCamelCase:					List[int]      ,		_UpperCamelCase:					int      ,		_UpperCamelCase:					int      ,		_UpperCamelCase:					Set[int]					)		->				Tuple[Set[int], torch.LongTensor]:

     """simple docstring"""




     _snake_case							=   torch.ones(_UpperCamelCase      ,		_UpperCamelCase					)
     _snake_case							=   set(_UpperCamelCase					) - 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
          _snake_case							=   head - sum(1 if h < head else 0 for h in already_pruned_heads					)
          _snake_case							=   0
     _snake_case							=   mask.view(-1					).contiguous().eq(1					)
     _snake_case							=   torch.arange(len(_UpperCamelCase					)					)[mask].long()
     return heads, index





def 			__a			(     *_UpperCamelCase:					Union[torch.Tensor, List[torch.Tensor]]      ,		_UpperCamelCase:					Optional[str] = None					)		->				Tuple[torch.Tensor, ...]:

     """simple docstring"""




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





def 			__a			(     _UpperCamelCase:					torch.Tensor					)		->				Tuple[torch.device, int, int]:

     """simple docstring"""




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








'''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      _a       (		unittest.TestCase							):






     def __init__(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE=13      ,_SCREAMING_SNAKE_CASE=7      ,_SCREAMING_SNAKE_CASE=True      ,_SCREAMING_SNAKE_CASE=True      ,_SCREAMING_SNAKE_CASE=True      ,_SCREAMING_SNAKE_CASE=True      ,_SCREAMING_SNAKE_CASE=99      ,_SCREAMING_SNAKE_CASE=24      ,_SCREAMING_SNAKE_CASE=32      ,_SCREAMING_SNAKE_CASE=5      ,_SCREAMING_SNAKE_CASE=4      ,_SCREAMING_SNAKE_CASE=37      ,_SCREAMING_SNAKE_CASE="gelu"      ,_SCREAMING_SNAKE_CASE=0.1      ,_SCREAMING_SNAKE_CASE=0.1      ,_SCREAMING_SNAKE_CASE=512      ,_SCREAMING_SNAKE_CASE=16      ,_SCREAMING_SNAKE_CASE=2      ,_SCREAMING_SNAKE_CASE=0.0_2      ,_SCREAMING_SNAKE_CASE=4      ,)							->    List[str]:
          _snake_case							=   parent
          _snake_case							=   batch_size
          _snake_case							=   seq_length
          _snake_case							=   is_training
          _snake_case							=   use_attention_mask
          _snake_case							=   use_token_type_ids
          _snake_case							=   use_labels
          _snake_case							=   vocab_size
          _snake_case							=   hidden_size
          _snake_case							=   embedding_size
          _snake_case							=   num_hidden_layers
          _snake_case							=   num_attention_heads
          _snake_case							=   intermediate_size
          _snake_case							=   hidden_act
          _snake_case							=   hidden_dropout_prob
          _snake_case							=   attention_probs_dropout_prob
          _snake_case							=   max_position_embeddings
          _snake_case							=   type_vocab_size
          _snake_case							=   type_sequence_label_size
          _snake_case							=   initializer_range
          _snake_case							=   num_choices






     def        _lowercase			(				self					)							->    int:
          _snake_case							=   ids_tensor([self.batch_size, self.seq_length]      ,self.vocab_size					)

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

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

          _snake_case							=   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        _lowercase			(				self					)							->    List[str]:
          _snake_case							=   self.prepare_config_and_inputs()
          _snake_case							,      _snake_case							,      _snake_case							,      _snake_case							=   config_and_inputs
          _snake_case							=   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
          return config, inputs_dict







@require_flax
class      _a       (		__lowerCAmelCase      ,      unittest.TestCase							):
     SCREAMING_SNAKE_CASE_							:					List[str]					       =		True

     SCREAMING_SNAKE_CASE_							:					Any					       =		(
         (
             FlaxElectraModel,
             FlaxElectraForCausalLM,
             FlaxElectraForMaskedLM,
             FlaxElectraForPreTraining,
             FlaxElectraForTokenClassification,
             FlaxElectraForQuestionAnswering,
             FlaxElectraForMultipleChoice,
             FlaxElectraForSequenceClassification,
         )
         if is_flax_available()
         else ()
     )






     def        _lowercase			(				self					)							->    List[Any]:
          _snake_case							=   FlaxElectraModelTester(self					)






     @slow
     def        _lowercase			(				self					)							->    List[Any]:
          for model_class_name in self.all_model_classes:
               if model_class_name == FlaxElectraForMaskedLM:
                    _snake_case							=   model_class_name.from_pretrained("google/electra-small-generator"					)
               else:
                    _snake_case							=   model_class_name.from_pretrained("google/electra-small-discriminator"					)
               _snake_case							=   model(np.ones((1, 1)					)					)
               self.assertIsNotNone(_SCREAMING_SNAKE_CASE					)


