



from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


_lowercase     							=						logging.get_logger(__name__)

# General docstring
_lowercase     							=						'''RegNetConfig'''

# Base docstring
_lowercase     							=						'''facebook/regnet-y-040'''
_lowercase     							=						[1, 10_88, 7, 7]

# Image classification docstring
_lowercase     							=						'''facebook/regnet-y-040'''
_lowercase     							=						'''tabby, tabby cat'''

_lowercase     							=						[
    '''facebook/regnet-y-040''',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]
class 				__A					( tf.keras.layers.Layer      ):





			def __init__(self							,    __magic_name__							,    __magic_name__ = 3							,    __magic_name__ = 1							,    __magic_name__ = 1							,    __magic_name__ = "relu"							,    **__magic_name__							,    ):
						super().__init__(**__magic_name__				)
						# The padding and conv has been verified in
						# https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
						lowerCamelCase__ :      Optional[Any]								=		tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2				)
						lowerCamelCase__ :      List[str]								=		tf.keras.layers.ConvaD(
						    filters=__magic_name__							,    kernel_size=__magic_name__							,    strides=__magic_name__							,    padding="""VALID"""							,    groups=__magic_name__							,    use_bias=__magic_name__							,    name="""convolution"""							,    )
						lowerCamelCase__ :      int								=		tf.keras.layers.BatchNormalization(epsilon=1E-5							,    momentum=0.9							,    name="""normalization"""				)
						lowerCamelCase__ :      Union[str, Any]								=		ACTaFN[activation] if activation is not None else tf.identity





			def        _snake_case    (self							,    __magic_name__				):
						lowerCamelCase__ :      Optional[Any]								=		self.convolution(self.padding(__magic_name__				)				)
						lowerCamelCase__ :      Optional[Any]								=		self.normalization(__magic_name__				)
						lowerCamelCase__ :      str								=		self.activation(__magic_name__				)
						return hidden_state







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





			def __init__(self							,    __magic_name__							,    **__magic_name__				):
						super().__init__(**__magic_name__				)
						lowerCamelCase__ :      int								=		config.num_channels
						lowerCamelCase__ :      Tuple								=		TFRegNetConvLayer(
						    out_channels=config.embedding_size							,    kernel_size=3							,    stride=2							,    activation=config.hidden_act							,    name="""embedder"""							,    )





			def        _snake_case    (self							,    __magic_name__				):
						lowerCamelCase__ :      Dict								=		shape_list(__magic_name__				)[1]
						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."""				)

						# When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
						# So change the input format from `NCHW` to `NHWC`.
						# shape = (batch_size, in_height, in_width, in_channels=num_channels)
						lowerCamelCase__ :      List[str]								=		tf.transpose(__magic_name__							,    perm=(0, 2, 3, 1)				)
						lowerCamelCase__ :      str								=		self.embedder(__magic_name__				)
						return hidden_state







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





			def __init__(self							,    __magic_name__							,    __magic_name__ = 2							,    **__magic_name__				):
						super().__init__(**__magic_name__				)
						lowerCamelCase__ :      int								=		tf.keras.layers.ConvaD(
						    filters=__magic_name__							,    kernel_size=1							,    strides=__magic_name__							,    use_bias=__magic_name__							,    name="""convolution"""				)
						lowerCamelCase__ :      Optional[int]								=		tf.keras.layers.BatchNormalization(epsilon=1E-5							,    momentum=0.9							,    name="""normalization"""				)





			def        _snake_case    (self							,    __magic_name__							,    __magic_name__ = False				):
						return self.normalization(self.convolution(__magic_name__				)							,    training=__magic_name__				)







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





			def __init__(self							,    __magic_name__							,    __magic_name__							,    **__magic_name__				):
						super().__init__(**__magic_name__				)
						lowerCamelCase__ :      List[str]								=		tf.keras.layers.GlobalAveragePoolingaD(keepdims=__magic_name__							,    name="""pooler"""				)
						lowerCamelCase__ :      Optional[int]								=		[
						    tf.keras.layers.ConvaD(filters=__magic_name__							,    kernel_size=1							,    activation="""relu"""							,    name="""attention.0"""				),
						    tf.keras.layers.ConvaD(filters=__magic_name__							,    kernel_size=1							,    activation="""sigmoid"""							,    name="""attention.2"""				),
						]





			def        _snake_case    (self							,    __magic_name__				):
						# [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
						lowerCamelCase__ :      Union[str, Any]								=		self.pooler(__magic_name__				)
						for layer_module in self.attention:
									lowerCamelCase__ :      Optional[Any]								=		layer_module(__magic_name__				)
						lowerCamelCase__ :      Union[str, Any]								=		hidden_state * pooled
						return hidden_state







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





			def __init__(self							,    __magic_name__							,    __magic_name__							,    __magic_name__							,    __magic_name__ = 1							,    **__magic_name__				):
						super().__init__(**__magic_name__				)
						lowerCamelCase__ :      str								=		in_channels != out_channels or stride != 1
						lowerCamelCase__ :      int								=		max(1							,    out_channels // config.groups_width				)
						lowerCamelCase__ :      int								=		(
						    TFRegNetShortCut(__magic_name__							,    stride=__magic_name__							,    name="""shortcut"""				)
						    if should_apply_shortcut
						    else tf.keras.layers.Activation("""linear"""							,    name="""shortcut"""				)
						)
						# `self.layers` instead of `self.layer` because that is a reserved argument.
						lowerCamelCase__ :      Optional[int]								=		[
						    TFRegNetConvLayer(__magic_name__							,    kernel_size=1							,    activation=config.hidden_act							,    name="""layer.0"""				),
						    TFRegNetConvLayer(
						        __magic_name__							,    stride=__magic_name__							,    groups=__magic_name__							,    activation=config.hidden_act							,    name="""layer.1"""				),
						    TFRegNetConvLayer(__magic_name__							,    kernel_size=1							,    activation=__magic_name__							,    name="""layer.2"""				),
						]
						lowerCamelCase__ :      Union[str, Any]								=		ACTaFN[config.hidden_act]





			def        _snake_case    (self							,    __magic_name__				):
						lowerCamelCase__ :      int								=		hidden_state
						for layer_module in self.layers:
									lowerCamelCase__ :      List[Any]								=		layer_module(__magic_name__				)
						lowerCamelCase__ :      str								=		self.shortcut(__magic_name__				)
						hidden_state += residual
						lowerCamelCase__ :      str								=		self.activation(__magic_name__				)
						return hidden_state







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





			def __init__(self							,    __magic_name__							,    __magic_name__							,    __magic_name__							,    __magic_name__ = 1							,    **__magic_name__				):
						super().__init__(**__magic_name__				)
						lowerCamelCase__ :      Union[str, Any]								=		in_channels != out_channels or stride != 1
						lowerCamelCase__ :      Optional[int]								=		max(1							,    out_channels // config.groups_width				)
						lowerCamelCase__ :      str								=		(
						    TFRegNetShortCut(__magic_name__							,    stride=__magic_name__							,    name="""shortcut"""				)
						    if should_apply_shortcut
						    else tf.keras.layers.Activation("""linear"""							,    name="""shortcut"""				)
						)
						lowerCamelCase__ :      Tuple								=		[
						    TFRegNetConvLayer(__magic_name__							,    kernel_size=1							,    activation=config.hidden_act							,    name="""layer.0"""				),
						    TFRegNetConvLayer(
						        __magic_name__							,    stride=__magic_name__							,    groups=__magic_name__							,    activation=config.hidden_act							,    name="""layer.1"""				),
						    TFRegNetSELayer(__magic_name__							,    reduced_channels=int(round(in_channels / 4				)				)							,    name="""layer.2"""				),
						    TFRegNetConvLayer(__magic_name__							,    kernel_size=1							,    activation=__magic_name__							,    name="""layer.3"""				),
						]
						lowerCamelCase__ :      str								=		ACTaFN[config.hidden_act]





			def        _snake_case    (self							,    __magic_name__				):
						lowerCamelCase__ :      Any								=		hidden_state
						for layer_module in self.layers:
									lowerCamelCase__ :      str								=		layer_module(__magic_name__				)
						lowerCamelCase__ :      Any								=		self.shortcut(__magic_name__				)
						hidden_state += residual
						lowerCamelCase__ :      List[str]								=		self.activation(__magic_name__				)
						return hidden_state







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





			def __init__(self							,    __magic_name__							,    __magic_name__							,    __magic_name__							,    __magic_name__ = 2							,    __magic_name__ = 2							,    **__magic_name__				):
						super().__init__(**__magic_name__				)

						lowerCamelCase__ :      Dict								=		TFRegNetXLayer if config.layer_type == """x""" else TFRegNetYLayer
						lowerCamelCase__ :      Union[str, Any]								=		[
						    # downsampling is done in the first layer with stride of 2
						    layer(__magic_name__							,    __magic_name__							,    __magic_name__							,    stride=__magic_name__							,    name="""layers.0"""				),
						    *[layer(__magic_name__							,    __magic_name__							,    __magic_name__							,    name=f"layers.{i+1}"				) for i in range(depth - 1				)],
						]





			def        _snake_case    (self							,    __magic_name__				):
						for layer_module in self.layers:
									lowerCamelCase__ :      Tuple								=		layer_module(__magic_name__				)
						return hidden_state







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





			def __init__(self							,    __magic_name__							,    **__magic_name__				):
						super().__init__(**__magic_name__				)
						lowerCamelCase__ :      Tuple								=		[]
						# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
						self.stages.append(
						    TFRegNetStage(
						        __magic_name__							,    config.embedding_size							,    config.hidden_sizes[0]							,    stride=2 if config.downsample_in_first_stage else 1							,    depth=config.depths[0]							,    name="""stages.0"""							,    )				)
						lowerCamelCase__ :      Optional[int]								=		zip(config.hidden_sizes							,    config.hidden_sizes[1:]				)
						for i, ((in_channels, out_channels), depth) in enumerate(zip(__magic_name__							,    config.depths[1:]				)				):
									self.stages.append(TFRegNetStage(__magic_name__							,    __magic_name__							,    __magic_name__							,    depth=__magic_name__							,    name=f"stages.{i+1}"				)				)





			def        _snake_case    (self							,    __magic_name__							,    __magic_name__ = False							,    __magic_name__ = True				):
						lowerCamelCase__ :      int								=		() if output_hidden_states else None

						for stage_module in self.stages:
									if output_hidden_states:
												lowerCamelCase__ :      List[Any]								=		hidden_states + (hidden_state,)

									lowerCamelCase__ :      Optional[Any]								=		stage_module(__magic_name__				)

						if output_hidden_states:
									lowerCamelCase__ :      Optional[Any]								=		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=__magic_name__							,    hidden_states=__magic_name__				)







@keras_serializable
class 				__A					( tf.keras.layers.Layer      ):
			UpperCamelCase					:Optional[Any]     =			RegNetConfig





			def __init__(self							,    __magic_name__							,    **__magic_name__				):
						super().__init__(**__magic_name__				)
						lowerCamelCase__ :      str								=		config
						lowerCamelCase__ :      Any								=		TFRegNetEmbeddings(__magic_name__							,    name="""embedder"""				)
						lowerCamelCase__ :      Optional[int]								=		TFRegNetEncoder(__magic_name__							,    name="""encoder"""				)
						lowerCamelCase__ :      Optional[int]								=		tf.keras.layers.GlobalAveragePoolingaD(keepdims=__magic_name__							,    name="""pooler"""				)





			@unpack_inputs
			def        _snake_case    (self							,    __magic_name__							,    __magic_name__ = None							,    __magic_name__ = None							,    __magic_name__ = False							,    ):
						lowerCamelCase__ :      Tuple								=		(
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						lowerCamelCase__ :      str								=		return_dict if return_dict is not None else self.config.use_return_dict

						lowerCamelCase__ :      str								=		self.embedder(__magic_name__							,    training=__magic_name__				)

						lowerCamelCase__ :      Optional[int]								=		self.encoder(
						    __magic_name__							,    output_hidden_states=__magic_name__							,    return_dict=__magic_name__							,    training=__magic_name__				)

						lowerCamelCase__ :      Optional[int]								=		encoder_outputs[0]
						lowerCamelCase__ :      Tuple								=		self.pooler(__magic_name__				)

						# Change to NCHW output format have uniformity in the modules
						lowerCamelCase__ :      Union[str, Any]								=		tf.transpose(__magic_name__							,    perm=(0, 3, 1, 2)				)
						lowerCamelCase__ :      List[Any]								=		tf.transpose(__magic_name__							,    perm=(0, 3, 1, 2)				)

						# Change the other hidden state outputs to NCHW as well
						if output_hidden_states:
									lowerCamelCase__ :      Optional[int]								=		tuple([tf.transpose(__magic_name__							,    perm=(0, 3, 1, 2)				) for h in encoder_outputs[1]]				)

						if not return_dict:
									return (last_hidden_state, pooled_output) + encoder_outputs[1:]

						return TFBaseModelOutputWithPoolingAndNoAttention(
						    last_hidden_state=__magic_name__							,    pooler_output=__magic_name__							,    hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states							,    )







class 				__A					( A_      ):
			UpperCamelCase					:Optional[Any]     =			RegNetConfig
			UpperCamelCase					:int     =			'''regnet'''
			UpperCamelCase					:Optional[Any]     =			'''pixel_values'''





			@property
			def        _snake_case    (self				):
						return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)							,    dtype=tf.floataa				)}







_lowercase     							=						r'''
    Parameters:
    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.
        config ([`RegNetConfig`]): 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.
'''

_lowercase     							=						r'''
    Args:
        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConveNextImageProcessor.__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.
'''
@add_start_docstrings(
    '''The bare RegNet model outputting raw features without any specific head on top.'''					, A_					, )
class 				__A					( A_      ):





			def __init__(self							,    __magic_name__							,    *__magic_name__							,    **__magic_name__				):
						super().__init__(__magic_name__							,    *__magic_name__							,    **__magic_name__				)
						lowerCamelCase__ :      int								=		TFRegNetMainLayer(__magic_name__							,    name="""regnet"""				)





			@unpack_inputs
			@add_start_docstrings_to_model_forward(__magic_name__				)
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC							,    output_type=__magic_name__							,    config_class=_CONFIG_FOR_DOC							,    modality="""vision"""							,    expected_output=_EXPECTED_OUTPUT_SHAPE							,    )
			def        _snake_case    (self							,    __magic_name__							,    __magic_name__ = None							,    __magic_name__ = None							,    __magic_name__=False							,    ):
						lowerCamelCase__ :      Optional[int]								=		(
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						lowerCamelCase__ :      Any								=		return_dict if return_dict is not None else self.config.use_return_dict

						lowerCamelCase__ :      Optional[int]								=		self.regnet(
						    pixel_values=__magic_name__							,    output_hidden_states=__magic_name__							,    return_dict=__magic_name__							,    training=__magic_name__							,    )
						if not return_dict:
									return (outputs[0],) + outputs[1:]

						return TFBaseModelOutputWithPoolingAndNoAttention(
						    last_hidden_state=outputs.last_hidden_state							,    pooler_output=outputs.pooler_output							,    hidden_states=outputs.hidden_states							,    )







@add_start_docstrings(
    '''
    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    '''					, A_					, )
class 				__A					( A_					, A_      ):





			def __init__(self							,    __magic_name__							,    *__magic_name__							,    **__magic_name__				):
						super().__init__(__magic_name__							,    *__magic_name__							,    **__magic_name__				)
						lowerCamelCase__ :      Any								=		config.num_labels
						lowerCamelCase__ :      Optional[int]								=		TFRegNetMainLayer(__magic_name__							,    name="""regnet"""				)
						# classification head
						lowerCamelCase__ :      Dict								=		[
						    tf.keras.layers.Flatten(),
						    tf.keras.layers.Dense(config.num_labels							,    name="""classifier.1"""				) if config.num_labels > 0 else tf.identity,
						]





			@unpack_inputs
			@add_start_docstrings_to_model_forward(__magic_name__				)
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT							,    output_type=__magic_name__							,    config_class=_CONFIG_FOR_DOC							,    expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							,    )
			def        _snake_case    (self							,    __magic_name__ = None							,    __magic_name__ = None							,    __magic_name__ = None							,    __magic_name__ = None							,    __magic_name__=False							,    ):
						lowerCamelCase__ :      Union[str, Any]								=		(
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						lowerCamelCase__ :      List[str]								=		return_dict if return_dict is not None else self.config.use_return_dict

						lowerCamelCase__ :      List[str]								=		self.regnet(
						    __magic_name__							,    output_hidden_states=__magic_name__							,    return_dict=__magic_name__							,    training=__magic_name__				)

						lowerCamelCase__ :      Optional[int]								=		outputs.pooler_output if return_dict else outputs[1]

						lowerCamelCase__ :      str								=		self.classifier[0](__magic_name__				)
						lowerCamelCase__ :      Optional[int]								=		self.classifier[1](__magic_name__				)

						lowerCamelCase__ :      str								=		None if labels is None else self.hf_compute_loss(labels=__magic_name__							,    logits=__magic_name__				)

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

						return TFSequenceClassifierOutput(loss=__magic_name__							,    logits=__magic_name__							,    hidden_states=outputs.hidden_states				)








from ... import PretrainedConfig


_lowercase     							=						{
    '''sijunhe/nezha-cn-base''': '''https://huggingface.co/sijunhe/nezha-cn-base/resolve/main/config.json''',
}
class 				__A					( A_      ):
			UpperCamelCase					:Optional[Any]     =			NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP
			UpperCamelCase					:int     =			'''nezha'''





			def __init__(self							,    __magic_name__=21128							,    __magic_name__=768							,    __magic_name__=12							,    __magic_name__=12							,    __magic_name__=3072							,    __magic_name__="gelu"							,    __magic_name__=0.1							,    __magic_name__=0.1							,    __magic_name__=512							,    __magic_name__=64							,    __magic_name__=2							,    __magic_name__=0.02							,    __magic_name__=1E-12							,    __magic_name__=0.1							,    __magic_name__=0							,    __magic_name__=2							,    __magic_name__=3							,    __magic_name__=True							,    **__magic_name__							,    ):
						super().__init__(pad_token_id=__magic_name__							,    bos_token_id=__magic_name__							,    eos_token_id=__magic_name__							,    **__magic_name__				)

						lowerCamelCase__ :      List[str]								=		vocab_size
						lowerCamelCase__ :      int								=		hidden_size
						lowerCamelCase__ :      int								=		num_hidden_layers
						lowerCamelCase__ :      Optional[int]								=		num_attention_heads
						lowerCamelCase__ :      Dict								=		hidden_act
						lowerCamelCase__ :      str								=		intermediate_size
						lowerCamelCase__ :      Optional[Any]								=		hidden_dropout_prob
						lowerCamelCase__ :      Dict								=		attention_probs_dropout_prob
						lowerCamelCase__ :      str								=		max_position_embeddings
						lowerCamelCase__ :      Optional[Any]								=		max_relative_position
						lowerCamelCase__ :      Optional[int]								=		type_vocab_size
						lowerCamelCase__ :      Optional[Any]								=		initializer_range
						lowerCamelCase__ :      Any								=		layer_norm_eps
						lowerCamelCase__ :      Union[str, Any]								=		classifier_dropout
						lowerCamelCase__ :      Union[str, Any]								=		use_cache



