


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







						'''simple docstring'''




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




						def     lowercase_			(					self				)       ->     Optional[Any]:
										lowerCAmelCase_			:   Optional[int]     					=					ids_tensor([self.batch_size, self.seq_length] ,    self.vocab_size				)

										lowerCAmelCase_			:   int     					=					None
										if self.use_attention_mask:
														lowerCAmelCase_			:   List[Any]     					=					random_attention_mask([self.batch_size, self.seq_length]				)

										lowerCAmelCase_			:   List[Any]     					=					None
										if self.use_token_type_ids:
														lowerCAmelCase_			:   Optional[int]     					=					ids_tensor([self.batch_size, self.seq_length] ,    self.type_vocab_size				)

										lowerCAmelCase_			:   Optional[int]     					=					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				)       ->     str:
										lowerCAmelCase_			:   List[str]     					=					self.prepare_config_and_inputs()
										lowerCAmelCase_					,	lowerCAmelCase_					,	lowerCAmelCase_					,	lowerCAmelCase_			:   Tuple     					=					config_and_inputs
										lowerCAmelCase_			:   int     					=					{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
										return config, inputs_dict







@require_flax
class 					snake_case__(				UpperCAmelCase__,    unittest.TestCase				):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__   :	str       =				True

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




						def     lowercase_			(					self				)       ->     Optional[int]:
										lowerCAmelCase_			:   Dict     					=					FlaxElectraModelTester(self				)




						@slow
						def     lowercase_			(					self				)       ->     List[Any]:
										for model_class_name in self.all_model_classes:
														if model_class_name == FlaxElectraForMaskedLM:
																		lowerCAmelCase_			:   Dict     					=					model_class_name.from_pretrained('''google/electra-small-generator'''				)
														else:
																		lowerCAmelCase_			:   Optional[Any]     					=					model_class_name.from_pretrained('''google/electra-small-discriminator'''				)
														lowerCAmelCase_			:   str     					=					model(np.ones((1, 1)				)				)
														self.assertIsNotNone(__lowercase				)



import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
						import transformer_engine.pytorch as te



def       lowerCAmelCase			(		lowerCAmelCase_			, lowerCAmelCase_=True			, lowerCAmelCase_=True			, lowerCAmelCase_=True			)->			Union[str, Any]:
				if not is_fpa_available():
								raise ImportError('''Using `convert_model` requires transformer_engine to be installed.'''			)
				for name, module in model.named_children():
								if isinstance(lowerCAmelCase_			, nn.Linear			) and to_transformer_engine and _convert_linear:
												# Return early if the linear layer weights are not multiples of 16
												if any(p % 16 != 0 for p in module.weight.shape			):
																return
												lowerCAmelCase_			:   Optional[int]     					=					module.bias is not None
												lowerCAmelCase_			:   Any     					=					te.Linear(
												    module.in_features			, module.out_features			, bias=lowerCAmelCase_			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   Dict     					=					module.weight.data.clone()
												if has_bias:
																lowerCAmelCase_			:   Any     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								elif isinstance(lowerCAmelCase_			, nn.LayerNorm			) and to_transformer_engine and _convert_ln:
												lowerCAmelCase_			:   List[str]     					=					te.LayerNorm(module.normalized_shape[0]			, eps=module.eps			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   Union[str, Any]     					=					module.weight.data.clone()
												lowerCAmelCase_			:   List[Any]     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								elif isinstance(lowerCAmelCase_			, te.Linear			) and not to_transformer_engine and _convert_linear:
												lowerCAmelCase_			:   Dict     					=					module.bias is not None
												lowerCAmelCase_			:   Union[str, Any]     					=					nn.Linear(
												    module.in_features			, module.out_features			, bias=lowerCAmelCase_			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   Optional[Any]     					=					module.weight.data.clone()
												if has_bias:
																lowerCAmelCase_			:   Union[str, Any]     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								elif isinstance(lowerCAmelCase_			, te.LayerNorm			) and not to_transformer_engine and _convert_ln:
												lowerCAmelCase_			:   Union[str, Any]     					=					nn.LayerNorm(module.normalized_shape[0]			, eps=module.eps			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   List[str]     					=					module.weight.data.clone()
												lowerCAmelCase_			:   List[Any]     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								else:
												convert_model(
												    lowerCAmelCase_			, to_transformer_engine=lowerCAmelCase_			, _convert_linear=_convert_linear			, _convert_ln=_convert_ln			, )




def       lowerCAmelCase			(		lowerCAmelCase_			)->			str:
				if not is_fpa_available():
								raise ImportError('''Using `has_transformer_engine_layers` requires transformer_engine to be installed.'''			)
				for m in model.modules():
								if isinstance(lowerCAmelCase_			, (te.LayerNorm, te.Linear)			):
												return True
				return False