



import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp


def 						_UpperCAmelCase							( snake_case      ,		snake_case      ,		snake_case      ,		snake_case      ,		snake_case = 40_96     ):





					"""simple docstring"""

					_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase         =       key.shape[-3:]
					_lowerCAmelCase         =       value.shape[-1]
					_lowerCAmelCase         =       min(snake_case      ,		snake_case     )
					_lowerCAmelCase         =       query / jnp.sqrt(snake_case     )

					@functools.partial(jax.checkpoint      ,		prevent_cse=snake_case     )
					def summarize_chunk(snake_case      ,		snake_case      ,		snake_case     ):
										_lowerCAmelCase         =       jnp.einsum("""...qhd,...khd->...qhk"""      ,		snake_case      ,		snake_case      ,		precision=snake_case     )

										_lowerCAmelCase         =       jnp.max(snake_case      ,		axis=-1      ,		keepdims=snake_case     )
										_lowerCAmelCase         =       jax.lax.stop_gradient(snake_case     )
										_lowerCAmelCase         =       jnp.exp(attn_weights - max_score     )

										_lowerCAmelCase         =       jnp.einsum("""...vhf,...qhv->...qhf"""      ,		snake_case      ,		snake_case      ,		precision=snake_case     )
										_lowerCAmelCase         =       jnp.einsum("""...qhk->...qh"""      ,		snake_case     )

										return (exp_values, exp_weights.sum(axis=-1     ), max_score)

					def chunk_scanner(snake_case     ):
										# julienne key array
										_lowerCAmelCase         =       jax.lax.dynamic_slice(
										    operand=snake_case      ,		start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]      ,		slice_sizes=list(key.shape[:-3]     ) + [key_chunk_size, num_heads, k_features]      ,		)

										# julienne value array
										_lowerCAmelCase         =       jax.lax.dynamic_slice(
										    operand=snake_case      ,		start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]      ,		slice_sizes=list(value.shape[:-3]     ) + [key_chunk_size, num_heads, v_features]      ,		)

										return summarize_chunk(snake_case      ,		snake_case      ,		snake_case     )

					_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase         =       jax.lax.map(f=snake_case      ,		xs=jnp.arange(0      ,		snake_case      ,		snake_case     )     )

					_lowerCAmelCase         =       jnp.max(snake_case      ,		axis=0      ,		keepdims=snake_case     )
					_lowerCAmelCase         =       jnp.exp(chunk_max - global_max     )

					chunk_values *= jnp.expand_dims(snake_case      ,		axis=-1     )
					chunk_weights *= max_diffs

					_lowerCAmelCase         =       chunk_values.sum(axis=0     )
					_lowerCAmelCase         =       jnp.expand_dims(snake_case      ,		-1     ).sum(axis=0     )

					return all_values / all_weights


def 						_UpperCAmelCase							( snake_case      ,		snake_case      ,		snake_case      ,		snake_case=jax.lax.Precision.HIGHEST      ,		snake_case = 10_24      ,		snake_case = 40_96     ):





					"""simple docstring"""

					_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase         =       query.shape[-3:]

					def chunk_scanner(snake_case      ,		snake_case     ):
										# julienne query array
										_lowerCAmelCase         =       jax.lax.dynamic_slice(
										    operand=snake_case      ,		start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]      ,		slice_sizes=list(query.shape[:-3]     ) + [min(snake_case      ,		snake_case     ), num_heads, q_features]      ,		)

										return (
										    chunk_idx + query_chunk_size,  # unused ignore it
										    _query_chunk_attention(
										        query=snake_case      ,		key=snake_case      ,		value=snake_case      ,		precision=snake_case      ,		key_chunk_size=snake_case     ),
										)

					_lowerCAmelCase      ,      _lowerCAmelCase         =       jax.lax.scan(
					    f=snake_case      ,		init=0      ,		xs=snake_case      ,		length=math.ceil(num_q / query_chunk_size     )     )

					return jnp.concatenate(snake_case      ,		axis=-3     )  # fuse the chunked result back





class    __lowerCAmelCase					(	nn.Module			):
				__lowerCamelCase					=							42
				__lowerCamelCase					=							8
				__lowerCamelCase					=							64
				__lowerCamelCase					=							0.0
				__lowerCamelCase					=							False
				__lowerCamelCase					=							jnp.floataa




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       self.dim_head * self.heads
									_lowerCAmelCase         =       self.dim_head**-0.5

									# Weights were exported with old names {to_q, to_k, to_v, to_out}
									_lowerCAmelCase         =       nn.Dense(_snake_case		,							use_bias=_snake_case		,							dtype=self.dtype		,							name="""to_q"""						)
									_lowerCAmelCase         =       nn.Dense(_snake_case		,							use_bias=_snake_case		,							dtype=self.dtype		,							name="""to_k"""						)
									_lowerCAmelCase         =       nn.Dense(_snake_case		,							use_bias=_snake_case		,							dtype=self.dtype		,							name="""to_v"""						)

									_lowerCAmelCase         =       nn.Dense(self.query_dim		,							dtype=self.dtype		,							name="""to_out_0"""						)
									_lowerCAmelCase         =       nn.Dropout(rate=self.dropout						)




				def 				snake_case  (				self		,							_snake_case						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase         =       tensor.shape
									_lowerCAmelCase         =       self.heads
									_lowerCAmelCase         =       tensor.reshape(_snake_case		,							_snake_case		,							_snake_case		,							dim // head_size						)
									_lowerCAmelCase         =       jnp.transpose(_snake_case		,							(0, 2, 1, 3)						)
									_lowerCAmelCase         =       tensor.reshape(batch_size * head_size		,							_snake_case		,							dim // head_size						)
									return tensor




				def 				snake_case  (				self		,							_snake_case						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase         =       tensor.shape
									_lowerCAmelCase         =       self.heads
									_lowerCAmelCase         =       tensor.reshape(batch_size // head_size		,							_snake_case		,							_snake_case		,							_snake_case						)
									_lowerCAmelCase         =       jnp.transpose(_snake_case		,							(0, 2, 1, 3)						)
									_lowerCAmelCase         =       tensor.reshape(batch_size // head_size		,							_snake_case		,							dim * head_size						)
									return tensor




				def __call__(				self		,							_snake_case		,							_snake_case=None		,							_snake_case=True						):



									"""simple docstring"""





									_lowerCAmelCase         =       hidden_states if context is None else context

									_lowerCAmelCase         =       self.query(_snake_case						)
									_lowerCAmelCase         =       self.key(_snake_case						)
									_lowerCAmelCase         =       self.value(_snake_case						)

									_lowerCAmelCase         =       self.reshape_heads_to_batch_dim(_snake_case						)
									_lowerCAmelCase         =       self.reshape_heads_to_batch_dim(_snake_case						)
									_lowerCAmelCase         =       self.reshape_heads_to_batch_dim(_snake_case						)

									if self.use_memory_efficient_attention:
														_lowerCAmelCase         =       query_states.transpose(1		,							0		,							2						)
														_lowerCAmelCase         =       key_states.transpose(1		,							0		,							2						)
														_lowerCAmelCase         =       value_states.transpose(1		,							0		,							2						)

														# this if statement create a chunk size for each layer of the unet
														# the chunk size is equal to the query_length dimension of the deepest layer of the unet

														_lowerCAmelCase         =       query_states.shape[-3]
														if flatten_latent_dim % 64 == 0:
																			_lowerCAmelCase         =       int(flatten_latent_dim / 64						)
														elif flatten_latent_dim % 16 == 0:
																			_lowerCAmelCase         =       int(flatten_latent_dim / 16						)
														elif flatten_latent_dim % 4 == 0:
																			_lowerCAmelCase         =       int(flatten_latent_dim / 4						)
														else:
																			_lowerCAmelCase         =       int(_snake_case						)

														_lowerCAmelCase         =       jax_memory_efficient_attention(
														    _snake_case		,							_snake_case		,							_snake_case		,							query_chunk_size=_snake_case		,							key_chunk_size=4096 * 4						)

														_lowerCAmelCase         =       hidden_states.transpose(1		,							0		,							2						)
									else:
														# compute attentions
														_lowerCAmelCase         =       jnp.einsum("""b i d, b j d->b i j"""		,							_snake_case		,							_snake_case						)
														_lowerCAmelCase         =       attention_scores * self.scale
														_lowerCAmelCase         =       nn.softmax(_snake_case		,							axis=2						)

														# attend to values
														_lowerCAmelCase         =       jnp.einsum("""b i j, b j d -> b i d"""		,							_snake_case		,							_snake_case						)

									_lowerCAmelCase         =       self.reshape_batch_dim_to_heads(_snake_case						)
									_lowerCAmelCase         =       self.proj_attn(_snake_case						)
									return self.dropout_layer(_snake_case		,							deterministic=_snake_case						)





class    __lowerCAmelCase					(	nn.Module			):
				__lowerCamelCase					=							42
				__lowerCamelCase					=							42
				__lowerCamelCase					=							42
				__lowerCamelCase					=							0.0
				__lowerCamelCase					=							False
				__lowerCamelCase					=							jnp.floataa
				__lowerCamelCase					=							False




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       FlaxAttention(
									    self.dim		,							self.n_heads		,							self.d_head		,							self.dropout		,							self.use_memory_efficient_attention		,							dtype=self.dtype						)
									# cross attention
									_lowerCAmelCase         =       FlaxAttention(
									    self.dim		,							self.n_heads		,							self.d_head		,							self.dropout		,							self.use_memory_efficient_attention		,							dtype=self.dtype						)
									_lowerCAmelCase         =       FlaxFeedForward(dim=self.dim		,							dropout=self.dropout		,							dtype=self.dtype						)
									_lowerCAmelCase         =       nn.LayerNorm(epsilon=1e-5		,							dtype=self.dtype						)
									_lowerCAmelCase         =       nn.LayerNorm(epsilon=1e-5		,							dtype=self.dtype						)
									_lowerCAmelCase         =       nn.LayerNorm(epsilon=1e-5		,							dtype=self.dtype						)
									_lowerCAmelCase         =       nn.Dropout(rate=self.dropout						)




				def __call__(				self		,							_snake_case		,							_snake_case		,							_snake_case=True						):



									"""simple docstring"""





									_lowerCAmelCase         =       hidden_states
									if self.only_cross_attention:
														_lowerCAmelCase         =       self.attna(self.norma(_snake_case						)		,							_snake_case		,							deterministic=_snake_case						)
									else:
														_lowerCAmelCase         =       self.attna(self.norma(_snake_case						)		,							deterministic=_snake_case						)
									_lowerCAmelCase         =       hidden_states + residual

									# cross attention
									_lowerCAmelCase         =       hidden_states
									_lowerCAmelCase         =       self.attna(self.norma(_snake_case						)		,							_snake_case		,							deterministic=_snake_case						)
									_lowerCAmelCase         =       hidden_states + residual

									# feed forward
									_lowerCAmelCase         =       hidden_states
									_lowerCAmelCase         =       self.ff(self.norma(_snake_case						)		,							deterministic=_snake_case						)
									_lowerCAmelCase         =       hidden_states + residual

									return self.dropout_layer(_snake_case		,							deterministic=_snake_case						)





class    __lowerCAmelCase					(	nn.Module			):
				__lowerCamelCase					=							42
				__lowerCamelCase					=							42
				__lowerCamelCase					=							42
				__lowerCamelCase					=							1
				__lowerCamelCase					=							0.0
				__lowerCamelCase					=							False
				__lowerCamelCase					=							False
				__lowerCamelCase					=							jnp.floataa
				__lowerCamelCase					=							False




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       nn.GroupNorm(num_groups=32		,							epsilon=1e-5						)

									_lowerCAmelCase         =       self.n_heads * self.d_head
									if self.use_linear_projection:
														_lowerCAmelCase         =       nn.Dense(_snake_case		,							dtype=self.dtype						)
									else:
														_lowerCAmelCase         =       nn.Conv(
														    _snake_case		,							kernel_size=(1, 1)		,							strides=(1, 1)		,							padding="""VALID"""		,							dtype=self.dtype		,							)

									_lowerCAmelCase         =       [
									    FlaxBasicTransformerBlock(
									        _snake_case		,							self.n_heads		,							self.d_head		,							dropout=self.dropout		,							only_cross_attention=self.only_cross_attention		,							dtype=self.dtype		,							use_memory_efficient_attention=self.use_memory_efficient_attention		,							)
									    for _ in range(self.depth						)
									]

									if self.use_linear_projection:
														_lowerCAmelCase         =       nn.Dense(_snake_case		,							dtype=self.dtype						)
									else:
														_lowerCAmelCase         =       nn.Conv(
														    _snake_case		,							kernel_size=(1, 1)		,							strides=(1, 1)		,							padding="""VALID"""		,							dtype=self.dtype		,							)

									_lowerCAmelCase         =       nn.Dropout(rate=self.dropout						)




				def __call__(				self		,							_snake_case		,							_snake_case		,							_snake_case=True						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase      ,      _lowerCAmelCase         =       hidden_states.shape
									_lowerCAmelCase         =       hidden_states
									_lowerCAmelCase         =       self.norm(_snake_case						)
									if self.use_linear_projection:
														_lowerCAmelCase         =       hidden_states.reshape(_snake_case		,							height * width		,							_snake_case						)
														_lowerCAmelCase         =       self.proj_in(_snake_case						)
									else:
														_lowerCAmelCase         =       self.proj_in(_snake_case						)
														_lowerCAmelCase         =       hidden_states.reshape(_snake_case		,							height * width		,							_snake_case						)

									for transformer_block in self.transformer_blocks:
														_lowerCAmelCase         =       transformer_block(_snake_case		,							_snake_case		,							deterministic=_snake_case						)

									if self.use_linear_projection:
														_lowerCAmelCase         =       self.proj_out(_snake_case						)
														_lowerCAmelCase         =       hidden_states.reshape(_snake_case		,							_snake_case		,							_snake_case		,							_snake_case						)
									else:
														_lowerCAmelCase         =       hidden_states.reshape(_snake_case		,							_snake_case		,							_snake_case		,							_snake_case						)
														_lowerCAmelCase         =       self.proj_out(_snake_case						)

									_lowerCAmelCase         =       hidden_states + residual
									return self.dropout_layer(_snake_case		,							deterministic=_snake_case						)





class    __lowerCAmelCase					(	nn.Module			):
				__lowerCamelCase					=							42
				__lowerCamelCase					=							0.0
				__lowerCamelCase					=							jnp.floataa




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       FlaxGEGLU(self.dim		,							self.dropout		,							self.dtype						)
									_lowerCAmelCase         =       nn.Dense(self.dim		,							dtype=self.dtype						)




				def __call__(				self		,							_snake_case		,							_snake_case=True						):



									"""simple docstring"""





									_lowerCAmelCase         =       self.net_a(_snake_case		,							deterministic=_snake_case						)
									_lowerCAmelCase         =       self.net_a(_snake_case						)
									return hidden_states





class    __lowerCAmelCase					(	nn.Module			):
				__lowerCamelCase					=							42
				__lowerCamelCase					=							0.0
				__lowerCamelCase					=							jnp.floataa




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       self.dim * 4
									_lowerCAmelCase         =       nn.Dense(inner_dim * 2		,							dtype=self.dtype						)
									_lowerCAmelCase         =       nn.Dropout(rate=self.dropout						)




				def __call__(				self		,							_snake_case		,							_snake_case=True						):



									"""simple docstring"""





									_lowerCAmelCase         =       self.proj(_snake_case						)
									_lowerCAmelCase      ,      _lowerCAmelCase         =       jnp.split(_snake_case		,							2		,							axis=2						)
									return self.dropout_layer(hidden_linear * nn.gelu(_snake_case						)		,							deterministic=_snake_case						)










'''simple docstring'''






import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    MODEL_FOR_PRETRAINING_MAPPING,
	    ElectraForCausalLM,
	    ElectraForMaskedLM,
	    ElectraForMultipleChoice,
	    ElectraForPreTraining,
	    ElectraForQuestionAnswering,
	    ElectraForSequenceClassification,
	    ElectraForTokenClassification,
	    ElectraModel,
	)
	from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST



class  lowerCamelCase__				:


	"""simple docstring"""





	def __init__(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__=1_3   ,     UpperCAmelCase__=7   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=True   ,     UpperCAmelCase__=9_9   ,     UpperCAmelCase__=3_2   ,     UpperCAmelCase__=5   ,     UpperCAmelCase__=4   ,     UpperCAmelCase__=3_7   ,     UpperCAmelCase__="gelu"   ,     UpperCAmelCase__=0.1   ,     UpperCAmelCase__=0.1   ,     UpperCAmelCase__=5_1_2   ,     UpperCAmelCase__=1_6   ,     UpperCAmelCase__=2   ,     UpperCAmelCase__=0.0_2   ,     UpperCAmelCase__=3   ,     UpperCAmelCase__=4   ,     UpperCAmelCase__=None   ,     )       ->				Optional[Any]:
					_A :				Optional[int]				=  parent
					_A :				List[Any]				=  batch_size
					_A :				Optional[int]				=  seq_length
					_A :				Dict				=  is_training
					_A :				Union[str, Any]				=  use_input_mask
					_A :				Tuple				=  use_token_type_ids
					_A :				str				=  use_labels
					_A :				List[str]				=  vocab_size
					_A :				Tuple				=  hidden_size
					_A :				Tuple				=  num_hidden_layers
					_A :				List[str]				=  num_attention_heads
					_A :				Optional[Any]				=  intermediate_size
					_A :				str				=  hidden_act
					_A :				Optional[int]				=  hidden_dropout_prob
					_A :				int				=  attention_probs_dropout_prob
					_A :				Optional[Any]				=  max_position_embeddings
					_A :				List[Any]				=  type_vocab_size
					_A :				Tuple				=  type_sequence_label_size
					_A :				Any				=  initializer_range
					_A :				Any				=  num_labels
					_A :				str				=  num_choices
					_A :				List[str]				=  scope




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					_A :				List[Any]				=  ids_tensor([self.batch_size, self.seq_length]   ,     self.vocab_size			)

					_A :				Union[str, Any]				=  None
					if self.use_input_mask:
									_A :				Tuple				=  random_attention_mask([self.batch_size, self.seq_length]			)

					_A :				List[str]				=  None
					if self.use_token_type_ids:
									_A :				List[Any]				=  ids_tensor([self.batch_size, self.seq_length]   ,     self.type_vocab_size			)

					_A :				Tuple				=  None
					_A :				Dict				=  None
					_A :				int				=  None
					if self.use_labels:
									_A :				Any				=  ids_tensor([self.batch_size]   ,     self.type_sequence_label_size			)
									_A :				Dict				=  ids_tensor([self.batch_size, self.seq_length]   ,     self.num_labels			)
									_A :				int				=  ids_tensor([self.batch_size]   ,     self.num_choices			)
									_A :				List[str]				=  ids_tensor([self.batch_size, self.seq_length]   ,     1			)

					_A :				List[Any]				=  self.get_config()

					return (
					    config,
					    input_ids,
					    token_type_ids,
					    input_mask,
					    sequence_labels,
					    token_labels,
					    choice_labels,
					    fake_token_labels,
					)




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					return ElectraConfig(
					    vocab_size=self.vocab_size   ,     hidden_size=self.hidden_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   ,     is_decoder=UpperCAmelCase__   ,     initializer_range=self.initializer_range   ,     )




	def 				_lowerCamelCase				(					self			)       ->				int:
					(
					    (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       
					) :				List[str]				=  self.prepare_config_and_inputs()

					_A :				Tuple				=  True
					_A :				List[str]				=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
					_A :				str				=  ids_tensor([self.batch_size, self.seq_length]   ,     vocab_size=2			)

					return (
					    config,
					    input_ids,
					    token_type_ids,
					    input_mask,
					    sequence_labels,
					    token_labels,
					    choice_labels,
					    encoder_hidden_states,
					    encoder_attention_mask,
					)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Tuple:
					_A :				List[str]				=  ElectraModel(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Tuple				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__			)
					_A :				List[Any]				=  model(UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__			)
					_A :				Tuple				=  model(UpperCAmelCase__			)
					self.parent.assertEqual(result.last_hidden_state.shape   ,     (self.batch_size, self.seq_length, self.hidden_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[str]:
					_A :				Dict				=  True
					_A :				Any				=  ElectraModel(UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				List[str]				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     encoder_hidden_states=UpperCAmelCase__   ,     encoder_attention_mask=UpperCAmelCase__   ,     )
					_A :				Optional[int]				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     encoder_hidden_states=UpperCAmelCase__   ,     )
					_A :				List[Any]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__			)
					self.parent.assertEqual(result.last_hidden_state.shape   ,     (self.batch_size, self.seq_length, self.hidden_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[str]:
					_A :				Optional[int]				=  ElectraForMaskedLM(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				List[str]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length, self.vocab_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Union[str, Any]:
					_A :				List[Any]				=  ElectraForCausalLM(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Optional[int]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length, self.vocab_size)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[str]:
					_A :				str				=  self.num_labels
					_A :				str				=  ElectraForTokenClassification(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				str				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length, self.num_labels)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[Any]:
					_A :				Optional[int]				=  self.num_labels
					_A :				Optional[int]				=  ElectraForPreTraining(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Dict				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Tuple:
					_A :				Optional[Any]				=  self.num_labels
					_A :				Optional[Any]				=  ElectraForSequenceClassification(UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Any				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__			)
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.num_labels)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				List[Any]:
					_A :				Optional[int]				=  ElectraForQuestionAnswering(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Union[str, Any]				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     start_positions=UpperCAmelCase__   ,     end_positions=UpperCAmelCase__   ,     )
					self.parent.assertEqual(result.start_logits.shape   ,     (self.batch_size, self.seq_length)			)
					self.parent.assertEqual(result.end_logits.shape   ,     (self.batch_size, self.seq_length)			)




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     )       ->				Union[str, Any]:
					_A :				Union[str, Any]				=  self.num_choices
					_A :				Union[str, Any]				=  ElectraForMultipleChoice(config=UpperCAmelCase__			)
					model.to(UpperCAmelCase__			)
					model.eval()
					_A :				Dict				=  input_ids.unsqueeze(1			).expand(-1   ,     self.num_choices   ,     -1			).contiguous()
					_A :				str				=  token_type_ids.unsqueeze(1			).expand(-1   ,     self.num_choices   ,     -1			).contiguous()
					_A :				Dict				=  input_mask.unsqueeze(1			).expand(-1   ,     self.num_choices   ,     -1			).contiguous()
					_A :				str				=  model(
					    UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__   ,     token_type_ids=UpperCAmelCase__   ,     labels=UpperCAmelCase__   ,     )
					self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.num_choices)			)




	def 				_lowerCamelCase				(					self			)       ->				List[str]:
					_A :				Optional[Any]				=  self.prepare_config_and_inputs()
					(
					    (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       (
					    _A
					)	,       
					) :				int				=  config_and_inputs
					_A :				Dict				=  {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
					return config, inputs_dict



@require_torch
class  lowerCamelCase__				(						snake_case_  ,							snake_case_  ,							unittest.TestCase						):


	"""simple docstring"""





	__magic_name__							   =		(
	    (
	        ElectraModel,
	        ElectraForPreTraining,
	        ElectraForMaskedLM,
	        ElectraForCausalLM,
	        ElectraForMultipleChoice,
	        ElectraForTokenClassification,
	        ElectraForSequenceClassification,
	        ElectraForQuestionAnswering,
	    )
	    if is_torch_available()
	    else ()
	)
	__magic_name__							   =		(
	    {
	        """feature-extraction""": ElectraModel,
	        """fill-mask""": ElectraForMaskedLM,
	        """question-answering""": ElectraForQuestionAnswering,
	        """text-classification""": ElectraForSequenceClassification,
	        """text-generation""": ElectraForCausalLM,
	        """token-classification""": ElectraForTokenClassification,
	        """zero-shot""": ElectraForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	__magic_name__							   =		True




	def 				_lowerCamelCase				(					self   ,     UpperCAmelCase__   ,     UpperCAmelCase__   ,     UpperCAmelCase__=False			)       ->				Optional[Any]:
					_A :				List[str]				=  super()._prepare_for_class(UpperCAmelCase__   ,     UpperCAmelCase__   ,     return_labels=UpperCAmelCase__			)

					if return_labels:
									if model_class in get_values(UpperCAmelCase__			):
													_A :				Tuple				=  torch.zeros(
													    (self.model_tester.batch_size, self.model_tester.seq_length)   ,     dtype=torch.long   ,     device=UpperCAmelCase__			)
					return inputs_dict




	def 				_lowerCamelCase				(					self			)       ->				int:
					_A :				Optional[Any]				=  ElectraModelTester(self			)
					_A :				List[str]				=  ConfigTester(self   ,     config_class=UpperCAmelCase__   ,     hidden_size=3_7			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					self.config_tester.run_common_tests()




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					_A :				Union[str, Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_model(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[str]:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_electra_model_as_decoder(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				int				=  self.model_tester.prepare_config_and_inputs()
					for type in ["absolute", "relative_key", "relative_key_query"]:
									_A :				List[str]				=  type
									self.model_tester.create_and_check_electra_model(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_masked_lm(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				Any:
					_A :				Dict				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_token_classification(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				Union[str, Any]:
					_A :				List[str]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_pretraining(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				Tuple:
					_A :				Any				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_sequence_classification(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				int:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_question_answering(*UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				Optional[Any]				=  self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_electra_for_multiple_choice(*UpperCAmelCase__			)




	@slow
	def 				_lowerCamelCase				(					self			)       ->				List[str]:
					for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									_A :				List[str]				=  ElectraModel.from_pretrained(UpperCAmelCase__			)
									self.assertIsNotNone(UpperCAmelCase__			)




	def 				_lowerCamelCase				(					self			)       ->				List[Any]:
					_A :				str				=  self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_electra_for_causal_lm(*UpperCAmelCase__			)



@require_torch
class  lowerCamelCase__				(						unittest.TestCase						):


	"""simple docstring"""





	@slow
	def 				_lowerCamelCase				(					self			)       ->				int:
					_A :				Any				=  ElectraModel.from_pretrained('''google/electra-small-discriminator'''			)
					_A :				Union[str, Any]				=  torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]			)
					_A :				Union[str, Any]				=  torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]			)
					_A :				Optional[int]				=  model(UpperCAmelCase__   ,     attention_mask=UpperCAmelCase__			)[0]
					_A :				Tuple				=  torch.Size((1, 1_1, 2_5_6)			)
					self.assertEqual(output.shape   ,     UpperCAmelCase__			)
					_A :				Optional[Any]				=  torch.tensor(
					    [[[0.4_4_7_1, 0.6_8_2_1, -0.3_2_6_5], [0.4_6_2_7, 0.5_2_5_5, -0.3_6_6_8], [0.4_5_3_2, 0.3_3_1_3, -0.4_3_4_4]]]			)

					self.assertTrue(torch.allclose(output[:, 1:4, 1:4]   ,     UpperCAmelCase__   ,     atol=1e-4			)			)
