





'''simple docstring'''






from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFImageClassifierOutputWithNoAttention,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_resnet import ResNetConfig


A__       :     Optional[Any]          =logging.get_logger(__name__)

# General docstring
A__       :     Union[str, Any]          ="""ResNetConfig"""

# Base docstring
A__       :     Dict          ="""microsoft/resnet-50"""
A__       :     Tuple          =[1, 20_48, 7, 7]

# Image classification docstring
A__       :     Dict          ="""microsoft/resnet-50"""
A__       :     int          ="""tiger cat"""

A__       :     Optional[Any]          =[
    """microsoft/resnet-50""",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]




class 				UpperCAmelCase  (    tf.keras.layers.Layer     ):




						def __init__(    self    :				List[str]    ,    __snake_case    :				int    ,    __snake_case    :				int = 3    ,    __snake_case    :				int = 1    ,    __snake_case    :				str = "relu"    ,    **__snake_case    :				Tuple						)       ->   List[Any]:
												super().__init__(**__snake_case						)
												_lowerCAmelCase							  =							kernel_size // 2
												_lowerCAmelCase							  =							tf.keras.layers.ConvaD(
												    __snake_case    ,    kernel_size=__snake_case    ,    strides=__snake_case    ,    padding="""valid"""    ,    use_bias=__snake_case    ,    name="""convolution"""						)
												# Use same default momentum and epsilon as PyTorch equivalent
												_lowerCAmelCase							  =							tf.keras.layers.BatchNormalization(epsilon=1E-5    ,    momentum=0.9    ,    name="""normalization"""						)
												_lowerCAmelCase							  =							ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("""linear"""						)




						def        lowercase__					(    self    :				int    ,    __snake_case    :				tf.Tensor						)       ->   Dict:
												# Pad to match that done in the PyTorch Conv2D model
												_lowerCAmelCase							  =							(self.pad_value, self.pad_value)
												_lowerCAmelCase							  =							tf.pad(__snake_case    ,    [(0, 0), height_pad, width_pad, (0, 0)]						)
												_lowerCAmelCase							  =							self.conv(__snake_case						)
												return hidden_state




						def        lowercase__					(    self    :				Tuple    ,    __snake_case    :				tf.Tensor    ,    __snake_case    :				bool = False						)       ->   Tuple:
												_lowerCAmelCase							  =							self.convolution(__snake_case						)
												_lowerCAmelCase							  =							self.normalization(__snake_case    ,    training=__snake_case						)
												_lowerCAmelCase							  =							self.activation(__snake_case						)
												return hidden_state




class 				UpperCAmelCase  (    tf.keras.layers.Layer     ):




						def __init__(    self    :				Dict    ,    __snake_case    :				ResNetConfig    ,    **__snake_case    :				Optional[Any]						)       ->   Tuple:
												super().__init__(**__snake_case						)
												_lowerCAmelCase							  =							TFResNetConvLayer(
												    config.embedding_size    ,    kernel_size=7    ,    stride=2    ,    activation=config.hidden_act    ,    name="""embedder"""    ,    )
												_lowerCAmelCase							  =							tf.keras.layers.MaxPoolaD(pool_size=3    ,    strides=2    ,    padding="""valid"""    ,    name="""pooler"""						)
												_lowerCAmelCase							  =							config.num_channels




						def        lowercase__					(    self    :				List[Any]    ,    __snake_case    :				tf.Tensor    ,    __snake_case    :				bool = False						)       ->   Union[str, Any]:
												_lowerCAmelCase							  =							shape_list(__snake_case						)
												if tf.executing_eagerly() and num_channels != self.num_channels:
																		raise ValueError(
																		    """Make sure that the channel dimension of the pixel values match with the one set in the configuration."""						)
												_lowerCAmelCase							  =							pixel_values
												_lowerCAmelCase							  =							self.embedder(__snake_case						)
												_lowerCAmelCase							  =							tf.pad(__snake_case    ,    [[0, 0], [1, 1], [1, 1], [0, 0]]						)
												_lowerCAmelCase							  =							self.pooler(__snake_case						)
												return hidden_state




class 				UpperCAmelCase  (    tf.keras.layers.Layer     ):




						def __init__(    self    :				Union[str, Any]    ,    __snake_case    :				int    ,    __snake_case    :				int = 2    ,    **__snake_case    :				List[Any]						)       ->   Dict:
												super().__init__(**__snake_case						)
												_lowerCAmelCase							  =							tf.keras.layers.ConvaD(
												    __snake_case    ,    kernel_size=1    ,    strides=__snake_case    ,    use_bias=__snake_case    ,    name="""convolution"""						)
												# Use same default momentum and epsilon as PyTorch equivalent
												_lowerCAmelCase							  =							tf.keras.layers.BatchNormalization(epsilon=1E-5    ,    momentum=0.9    ,    name="""normalization"""						)




						def        lowercase__					(    self    :				Tuple    ,    __snake_case    :				tf.Tensor    ,    __snake_case    :				bool = False						)       ->   Dict:
												_lowerCAmelCase							  =							x
												_lowerCAmelCase							  =							self.convolution(__snake_case						)
												_lowerCAmelCase							  =							self.normalization(__snake_case    ,    training=__snake_case						)
												return hidden_state




class 				UpperCAmelCase  (    tf.keras.layers.Layer     ):




						def __init__(    self    :				List[Any]    ,    __snake_case    :				int    ,    __snake_case    :				int    ,    __snake_case    :				int = 1    ,    __snake_case    :				str = "relu"    ,    **__snake_case    :				Optional[int]						)       ->   Optional[Any]:
												super().__init__(**__snake_case						)
												_lowerCAmelCase							  =							in_channels != out_channels or stride != 1
												_lowerCAmelCase							  =							TFResNetConvLayer(__snake_case    ,    stride=__snake_case    ,    name="""layer.0"""						)
												_lowerCAmelCase							  =							TFResNetConvLayer(__snake_case    ,    activation=__snake_case    ,    name="""layer.1"""						)
												_lowerCAmelCase							  =							(
												    TFResNetShortCut(__snake_case    ,    stride=__snake_case    ,    name="""shortcut"""						)
												    if should_apply_shortcut
												    else tf.keras.layers.Activation("""linear"""    ,    name="""shortcut"""						)
												)
												_lowerCAmelCase							  =							ACTaFN[activation]




						def        lowercase__					(    self    :				List[Any]    ,    __snake_case    :				tf.Tensor    ,    __snake_case    :				bool = False						)       ->   Union[str, Any]:
												_lowerCAmelCase							  =							hidden_state
												_lowerCAmelCase							  =							self.conva(__snake_case    ,    training=__snake_case						)
												_lowerCAmelCase							  =							self.conva(__snake_case    ,    training=__snake_case						)
												_lowerCAmelCase							  =							self.shortcut(__snake_case    ,    training=__snake_case						)
												hidden_state += residual
												_lowerCAmelCase							  =							self.activation(__snake_case						)
												return hidden_state




class 				UpperCAmelCase  (    tf.keras.layers.Layer     ):




						def __init__(    self    :				Optional[Any]    ,    __snake_case    :				int    ,    __snake_case    :				int    ,    __snake_case    :				int = 1    ,    __snake_case    :				str = "relu"    ,    __snake_case    :				int = 4    ,    **__snake_case    :				Optional[int]    ,    )       ->   Optional[int]:
												super().__init__(**__snake_case						)
												_lowerCAmelCase							  =							in_channels != out_channels or stride != 1
												_lowerCAmelCase							  =							out_channels // reduction
												_lowerCAmelCase							  =							TFResNetConvLayer(__snake_case    ,    kernel_size=1    ,    name="""layer.0"""						)
												_lowerCAmelCase							  =							TFResNetConvLayer(__snake_case    ,    stride=__snake_case    ,    name="""layer.1"""						)
												_lowerCAmelCase							  =							TFResNetConvLayer(__snake_case    ,    kernel_size=1    ,    activation=__snake_case    ,    name="""layer.2"""						)
												_lowerCAmelCase							  =							(
												    TFResNetShortCut(__snake_case    ,    stride=__snake_case    ,    name="""shortcut"""						)
												    if should_apply_shortcut
												    else tf.keras.layers.Activation("""linear"""    ,    name="""shortcut"""						)
												)
												_lowerCAmelCase							  =							ACTaFN[activation]




						def        lowercase__					(    self    :				str    ,    __snake_case    :				tf.Tensor    ,    __snake_case    :				bool = False						)       ->   int:
												_lowerCAmelCase							  =							hidden_state
												_lowerCAmelCase							  =							self.conva(__snake_case    ,    training=__snake_case						)
												_lowerCAmelCase							  =							self.conva(__snake_case    ,    training=__snake_case						)
												_lowerCAmelCase							  =							self.conva(__snake_case    ,    training=__snake_case						)
												_lowerCAmelCase							  =							self.shortcut(__snake_case    ,    training=__snake_case						)
												hidden_state += residual
												_lowerCAmelCase							  =							self.activation(__snake_case						)
												return hidden_state




class 				UpperCAmelCase  (    tf.keras.layers.Layer     ):




						def __init__(    self    :				int    ,    __snake_case    :				ResNetConfig    ,    __snake_case    :				int    ,    __snake_case    :				int    ,    __snake_case    :				int = 2    ,    __snake_case    :				int = 2    ,    **__snake_case    :				List[Any]						)       ->   Optional[int]:
												super().__init__(**__snake_case						)

												_lowerCAmelCase							  =							TFResNetBottleNeckLayer if config.layer_type == 'bottleneck' else TFResNetBasicLayer

												_lowerCAmelCase							  =							[layer(__snake_case    ,    __snake_case    ,    stride=__snake_case    ,    activation=config.hidden_act    ,    name="""layers.0"""						)]
												layers += [
												    layer(__snake_case    ,    __snake_case    ,    activation=config.hidden_act    ,    name=f"layers.{i + 1}"						)
												    for i in range(depth - 1						)
												]
												_lowerCAmelCase							  =							layers




						def        lowercase__					(    self    :				Any    ,    __snake_case    :				tf.Tensor    ,    __snake_case    :				bool = False						)       ->   Any:
												for layer in self.stage_layers:
																		_lowerCAmelCase							  =							layer(__snake_case    ,    training=__snake_case						)
												return hidden_state




class 				UpperCAmelCase  (    tf.keras.layers.Layer     ):




						def __init__(    self    :				List[Any]    ,    __snake_case    :				ResNetConfig    ,    **__snake_case    :				Optional[Any]						)       ->   List[str]:
												super().__init__(**__snake_case						)
												# based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
												_lowerCAmelCase							  =							[
												    TFResNetStage(
												        __snake_case    ,    config.embedding_size    ,    config.hidden_sizes[0]    ,    stride=2 if config.downsample_in_first_stage else 1    ,    depth=config.depths[0]    ,    name="""stages.0"""    ,    )
												]
												for i, (in_channels, out_channels, depth) in enumerate(
												    zip(config.hidden_sizes    ,    config.hidden_sizes[1:]    ,    config.depths[1:]						)						):
																		self.stages.append(TFResNetStage(__snake_case    ,    __snake_case    ,    __snake_case    ,    depth=__snake_case    ,    name=f"stages.{i + 1}"						)						)




						def        lowercase__					(    self    :				Optional[int]    ,    __snake_case    :				tf.Tensor    ,    __snake_case    :				bool = False    ,    __snake_case    :				bool = True    ,    __snake_case    :				bool = False    ,    )       ->   List[Any]:
												_lowerCAmelCase							  =							() if output_hidden_states else None

												for stage_module in self.stages:
																		if output_hidden_states:
																								_lowerCAmelCase							  =							hidden_states + (hidden_state,)

																		_lowerCAmelCase							  =							stage_module(__snake_case    ,    training=__snake_case						)

												if output_hidden_states:
																		_lowerCAmelCase							  =							hidden_states + (hidden_state,)

												if not return_dict:
																		return tuple(v for v in [hidden_state, hidden_states] if v is not None						)

												return TFBaseModelOutputWithNoAttention(last_hidden_state=__snake_case    ,    hidden_states=__snake_case						)




class 				UpperCAmelCase  (    snake_case__     ):
						_lowercase:    int		  =     ResNetConfig
						_lowercase:    Optional[Any]		  =     'resnet'
						_lowercase:    int		  =     'pixel_values'




						@property
						def        lowercase__					(    self    :				Tuple						)       ->   Optional[int]:
												return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_24, 2_24)    ,    dtype=tf.floataa						)}




A__       :     Union[str, Any]          =r"""
    This model is a TensorFlow
    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a
    regular TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
"""


A__       :     Optional[Any]          =r"""
    Args:
        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""




@keras_serializable
class 				UpperCAmelCase  (    tf.keras.layers.Layer     ):
						_lowercase:    int		  =     ResNetConfig




						def __init__(    self    :				Dict    ,    __snake_case    :				ResNetConfig    ,    **__snake_case    :				List[str]						)       ->   List[str]:
												super().__init__(**__snake_case						)
												_lowerCAmelCase							  =							config
												_lowerCAmelCase							  =							TFResNetEmbeddings(__snake_case    ,    name="""embedder"""						)
												_lowerCAmelCase							  =							TFResNetEncoder(__snake_case    ,    name="""encoder"""						)
												_lowerCAmelCase							  =							tf.keras.layers.GlobalAveragePoolingaD(keepdims=__snake_case						)




						@unpack_inputs
						def        lowercase__					(    self    :				List[str]    ,    __snake_case    :				tf.Tensor    ,    __snake_case    :				Optional[bool] = None    ,    __snake_case    :				Optional[bool] = None    ,    __snake_case    :				bool = False    ,    )       ->   List[str]:
												_lowerCAmelCase							  =							(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												_lowerCAmelCase							  =							return_dict if return_dict is not None else self.config.use_return_dict

												# TF 2.0 image layers can't use NCHW format when running on CPU.
												# We transpose to NHWC format and then transpose back after the full forward pass.
												# (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
												_lowerCAmelCase							  =							tf.transpose(__snake_case    ,    perm=[0, 2, 3, 1]						)
												_lowerCAmelCase							  =							self.embedder(__snake_case    ,    training=__snake_case						)

												_lowerCAmelCase							  =							self.encoder(
												    __snake_case    ,    output_hidden_states=__snake_case    ,    return_dict=__snake_case    ,    training=__snake_case						)

												_lowerCAmelCase							  =							encoder_outputs[0]

												_lowerCAmelCase							  =							self.pooler(__snake_case						)

												# Transpose all the outputs to the NCHW format
												# (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
												_lowerCAmelCase							  =							tf.transpose(__snake_case    ,    (0, 3, 1, 2)						)
												_lowerCAmelCase							  =							tf.transpose(__snake_case    ,    (0, 3, 1, 2)						)
												_lowerCAmelCase							  =							()
												for hidden_state in encoder_outputs[1:]:
																		_lowerCAmelCase							  =							hidden_states + tuple(tf.transpose(__snake_case    ,    (0, 3, 1, 2)						) for h in hidden_state						)

												if not return_dict:
																		return (last_hidden_state, pooled_output) + hidden_states

												_lowerCAmelCase							  =							hidden_states if output_hidden_states else None

												return TFBaseModelOutputWithPoolingAndNoAttention(
												    last_hidden_state=__snake_case    ,    pooler_output=__snake_case    ,    hidden_states=__snake_case    ,    )




@add_start_docstrings(
    '''The bare ResNet model outputting raw features without any specific head on top.'''							,	snake_case__							,	)
class 				UpperCAmelCase  (    snake_case__     ):




						def __init__(    self    :				Any    ,    __snake_case    :				ResNetConfig    ,    **__snake_case    :				int						)       ->   int:
												super().__init__(__snake_case    ,    **__snake_case						)
												_lowerCAmelCase							  =							TFResNetMainLayer(config=__snake_case    ,    name="""resnet"""						)




						@add_start_docstrings_to_model_forward(__snake_case						)
						@add_code_sample_docstrings(
						    checkpoint=_CHECKPOINT_FOR_DOC    ,    output_type=__snake_case    ,    config_class=_CONFIG_FOR_DOC    ,    modality="""vision"""    ,    expected_output=_EXPECTED_OUTPUT_SHAPE    ,    )
						@unpack_inputs
						def        lowercase__					(    self    :				List[str]    ,    __snake_case    :				tf.Tensor    ,    __snake_case    :				Optional[bool] = None    ,    __snake_case    :				Optional[bool] = None    ,    __snake_case    :				bool = False    ,    )       ->   Union[str, Any]:
												_lowerCAmelCase							  =							(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												_lowerCAmelCase							  =							return_dict if return_dict is not None else self.config.use_return_dict

												_lowerCAmelCase							  =							self.resnet(
												    pixel_values=__snake_case    ,    output_hidden_states=__snake_case    ,    return_dict=__snake_case    ,    training=__snake_case    ,    )
												return resnet_outputs




@add_start_docstrings(
    '''\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    '''							,	snake_case__							,	)
class 				UpperCAmelCase  (    snake_case__							,	snake_case__     ):




						def __init__(    self    :				Tuple    ,    __snake_case    :				ResNetConfig    ,    **__snake_case    :				Optional[Any]						)       ->   Any:
												super().__init__(__snake_case    ,    **__snake_case						)
												_lowerCAmelCase							  =							config.num_labels
												_lowerCAmelCase							  =							TFResNetMainLayer(__snake_case    ,    name="""resnet"""						)
												# classification head
												_lowerCAmelCase							  =							(
												    tf.keras.layers.Dense(config.num_labels    ,    name="""classifier.1"""						)
												    if config.num_labels > 0
												    else tf.keras.layers.Activation("""linear"""    ,    name="""classifier.1"""						)
												)




						def        lowercase__					(    self    :				List[str]    ,    __snake_case    :				tf.Tensor						)       ->   Any:
												_lowerCAmelCase							  =							tf.keras.layers.Flatten()(__snake_case						)
												_lowerCAmelCase							  =							self.classifier_layer(__snake_case						)
												return logits




						@add_start_docstrings_to_model_forward(__snake_case						)
						@add_code_sample_docstrings(
						    checkpoint=_IMAGE_CLASS_CHECKPOINT    ,    output_type=__snake_case    ,    config_class=_CONFIG_FOR_DOC    ,    expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT    ,    )
						@unpack_inputs
						def        lowercase__					(    self    :				Optional[int]    ,    __snake_case    :				tf.Tensor = None    ,    __snake_case    :				tf.Tensor = None    ,    __snake_case    :				bool = None    ,    __snake_case    :				bool = None    ,    __snake_case    :				bool = False    ,    )       ->   Optional[Any]:
												_lowerCAmelCase							  =							return_dict if return_dict is not None else self.config.use_return_dict

												_lowerCAmelCase							  =							self.resnet(
												    __snake_case    ,    output_hidden_states=__snake_case    ,    return_dict=__snake_case    ,    training=__snake_case						)

												_lowerCAmelCase							  =							outputs.pooler_output if return_dict else outputs[1]

												_lowerCAmelCase							  =							self.classifier(__snake_case						)

												_lowerCAmelCase							  =							None if labels is None else self.hf_compute_loss(__snake_case    ,    __snake_case						)

												if not return_dict:
																		_lowerCAmelCase							  =							(logits,) + outputs[2:]
																		return (loss,) + output if loss is not None else output

												return TFImageClassifierOutputWithNoAttention(loss=__snake_case    ,    logits=__snake_case    ,    hidden_states=outputs.hidden_states						)













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




      def __init__(    self    :				str    ,    __snake_case    :				str    ,    __snake_case    :				str=13    ,    __snake_case    :				Optional[Any]=7    ,    __snake_case    :				str=True    ,    __snake_case    :				Any=True    ,    __snake_case    :				List[str]=True    ,    __snake_case    :				Any=True    ,    __snake_case    :				List[str]=99    ,    __snake_case    :				Dict=24    ,    __snake_case    :				Any=32    ,    __snake_case    :				List[str]=5    ,    __snake_case    :				Any=4    ,    __snake_case    :				Optional[Any]=37    ,    __snake_case    :				Union[str, Any]="gelu"    ,    __snake_case    :				Dict=0.1    ,    __snake_case    :				int=0.1    ,    __snake_case    :				Dict=5_12    ,    __snake_case    :				List[Any]=16    ,    __snake_case    :				Dict=2    ,    __snake_case    :				str=0.02    ,    __snake_case    :				Any=4    ,    )       ->   Union[str, Any]:
            _lowerCAmelCase							  =							parent
            _lowerCAmelCase							  =							batch_size
            _lowerCAmelCase							  =							seq_length
            _lowerCAmelCase							  =							is_training
            _lowerCAmelCase							  =							use_attention_mask
            _lowerCAmelCase							  =							use_token_type_ids
            _lowerCAmelCase							  =							use_labels
            _lowerCAmelCase							  =							vocab_size
            _lowerCAmelCase							  =							hidden_size
            _lowerCAmelCase							  =							embedding_size
            _lowerCAmelCase							  =							num_hidden_layers
            _lowerCAmelCase							  =							num_attention_heads
            _lowerCAmelCase							  =							intermediate_size
            _lowerCAmelCase							  =							hidden_act
            _lowerCAmelCase							  =							hidden_dropout_prob
            _lowerCAmelCase							  =							attention_probs_dropout_prob
            _lowerCAmelCase							  =							max_position_embeddings
            _lowerCAmelCase							  =							type_vocab_size
            _lowerCAmelCase							  =							type_sequence_label_size
            _lowerCAmelCase							  =							initializer_range
            _lowerCAmelCase							  =							num_choices




      def        lowercase__					(    self    :				Optional[Any]						)       ->   Optional[Any]:
            _lowerCAmelCase							  =							ids_tensor([self.batch_size, self.seq_length]    ,    self.vocab_size						)

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

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

            _lowerCAmelCase							  =							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    :				Any						)       ->   Optional[int]:
            _lowerCAmelCase							  =							self.prepare_config_and_inputs()
            _lowerCAmelCase  ,							_lowerCAmelCase  ,							_lowerCAmelCase  ,							_lowerCAmelCase							  =							config_and_inputs
            _lowerCAmelCase							  =							{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
            return config, inputs_dict




@require_flax
class 				UpperCAmelCase  (    snake_case_							,	unittest.TestCase     ):
      _lowercase:    str		  =     True

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




      def        lowercase__					(    self    :				Tuple						)       ->   Optional[int]:
            _lowerCAmelCase							  =							FlaxElectraModelTester(self						)




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






