



'''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__			(lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	self						):
    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__			(lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_ = 0						):
    UpperCamelCase__	     =      index.to(layer.weight.device						)
    UpperCamelCase__	     =      layer.weight.index_select(lowerCAmelCase_       ,	lowerCAmelCase_						).clone().detach()
    if layer.bias is not None:
        if dim == 1:
            UpperCamelCase__	     =      layer.bias.clone().detach()
        else:
            UpperCamelCase__	     =      layer.bias[index].clone().detach()
    UpperCamelCase__	     =      list(layer.weight.size()						)
    UpperCamelCase__	     =      len(lowerCAmelCase_						)
    UpperCamelCase__	     =      nn.Linear(new_size[1]       ,	new_size[0]       ,	bias=layer.bias is not None						).to(layer.weight.device						)
    UpperCamelCase__	     =      False
    new_layer.weight.copy_(W.contiguous()						)
    UpperCamelCase__	     =      True
    if layer.bias is not None:
        UpperCamelCase__	     =      False
        new_layer.bias.copy_(b.contiguous()						)
        UpperCamelCase__	     =      True
    return new_layer

class    lowerCamelCase_				(						nn.Module	):






   '''simple docstring'''

   def __init__(	self  ,    __UpperCamelCase  ,    __UpperCamelCase			):
       super().__init__()
       UpperCamelCase__	     =      nf
       UpperCamelCase__	     =      nn.Parameter(torch.empty(__UpperCamelCase  ,    __UpperCamelCase			)			)
       UpperCamelCase__	     =      nn.Parameter(torch.zeros(__UpperCamelCase			)			)
       nn.init.normal_(self.weight  ,    std=0.02			)







   def 						_UpperCamelCase							(	self  ,    __UpperCamelCase			):
       UpperCamelCase__	     =      x.size()[:-1] + (self.nf,)
       UpperCamelCase__	     =      torch.addmm(self.bias  ,    x.view(-1  ,    x.size(-1			)			)  ,    self.weight			)
       UpperCamelCase__	     =      x.view(__UpperCamelCase			)
       return x





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





def     snake_case__			(lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_ = None						):
    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__			(lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	*lowerCAmelCase_						):
    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
    UpperCamelCase__	     =      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:
        UpperCamelCase__	     =      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}'						)

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

        # chunk input tensor into tuples
        UpperCamelCase__	     =      tuple(input_tensor.chunk(lowerCAmelCase_       ,	dim=lowerCAmelCase_						) for input_tensor in input_tensors						)
        # apply forward fn to every tuple
        UpperCamelCase__	     =      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__			(lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_						):
    UpperCamelCase__	     =      torch.ones(lowerCAmelCase_       ,	lowerCAmelCase_						)
    UpperCamelCase__	     =      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
        UpperCamelCase__	     =      head - sum(1 if h < head else 0 for h in already_pruned_heads						)
        UpperCamelCase__	     =      0
    UpperCamelCase__	     =      mask.view(-1						).contiguous().eq(1						)
    UpperCamelCase__	     =      torch.arange(len(lowerCAmelCase_						)						)[mask].long()
    return heads, index





def     snake_case__			(*lowerCAmelCase_       ,	lowerCAmelCase_ = None						):
    return torch.meshgrid(*lowerCAmelCase_       ,	indexing=lowerCAmelCase_						)





def     snake_case__			(lowerCAmelCase_						):
    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			)




