
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__		:Optional[Any]			  =   logging.get_logger(__name__)

# General docstring
lowercase__		:Dict			  =   "RegNetConfig"

# Base docstring
lowercase__		:List[str]			  =   "facebook/regnet-y-040"
lowercase__		:Optional[Any]			  =   [1, 1088, 7, 7]

# Image classification docstring
lowercase__		:int			  =   "facebook/regnet-y-040"
lowercase__		:Tuple			  =   "tabby, tabby cat"

lowercase__		:Dict			  =   [
    "facebook/regnet-y-040",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]







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




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





							def 			A__				(    self       ,A__):
												lowercase						     =						self.convolution(self.padding(A__))
												lowercase						     =						self.normalization(A__)
												lowercase						     =						self.activation(A__)
												return hidden_state







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




							def __init__(    self       ,A__       ,**A__):
												super().__init__(**A__)
												lowercase						     =						config.num_channels
												lowercase						     =						TFRegNetConvLayer(
												    out_channels=config.embedding_size       ,kernel_size=3       ,stride=2       ,activation=config.hidden_act       ,name='''embedder'''       ,)





							def 			A__				(    self       ,A__):
												lowercase						     =						shape_list(A__)[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)
												lowercase						     =						tf.transpose(A__       ,perm=(0, 2, 3, 1))
												lowercase						     =						self.embedder(A__)
												return hidden_state







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




							def __init__(    self       ,A__       ,A__ = 2       ,**A__):
												super().__init__(**A__)
												lowercase						     =						tf.keras.layers.ConvaD(
												    filters=A__       ,kernel_size=1       ,strides=A__       ,use_bias=A__       ,name='''convolution''')
												lowercase						     =						tf.keras.layers.BatchNormalization(epsilon=1E-5       ,momentum=0.9       ,name='''normalization''')





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







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




							def __init__(    self       ,A__       ,A__       ,**A__):
												super().__init__(**A__)
												lowercase						     =						tf.keras.layers.GlobalAveragePoolingaD(keepdims=A__       ,name='''pooler''')
												lowercase						     =						[
												    tf.keras.layers.ConvaD(filters=A__       ,kernel_size=1       ,activation='''relu'''       ,name='''attention.0'''),
												    tf.keras.layers.ConvaD(filters=A__       ,kernel_size=1       ,activation='''sigmoid'''       ,name='''attention.2'''),
												]





							def 			A__				(    self       ,A__):
												# [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
												lowercase						     =						self.pooler(A__)
												for layer_module in self.attention:
																	lowercase						     =						layer_module(A__)
												lowercase						     =						hidden_state * pooled
												return hidden_state







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




							def __init__(    self       ,A__       ,A__       ,A__       ,A__ = 1       ,**A__):
												super().__init__(**A__)
												lowercase						     =						in_channels != out_channels or stride != 1
												lowercase						     =						max(1       ,out_channels // config.groups_width)
												lowercase						     =						(
												    TFRegNetShortCut(A__       ,stride=A__       ,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.
												lowercase						     =						[
												    TFRegNetConvLayer(A__       ,kernel_size=1       ,activation=config.hidden_act       ,name='''layer.0'''),
												    TFRegNetConvLayer(
												        A__       ,stride=A__       ,groups=A__       ,activation=config.hidden_act       ,name='''layer.1'''),
												    TFRegNetConvLayer(A__       ,kernel_size=1       ,activation=A__       ,name='''layer.2'''),
												]
												lowercase						     =						ACTaFN[config.hidden_act]





							def 			A__				(    self       ,A__):
												lowercase						     =						hidden_state
												for layer_module in self.layers:
																	lowercase						     =						layer_module(A__)
												lowercase						     =						self.shortcut(A__)
												hidden_state += residual
												lowercase						     =						self.activation(A__)
												return hidden_state







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




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





							def 			A__				(    self       ,A__):
												lowercase						     =						hidden_state
												for layer_module in self.layers:
																	lowercase						     =						layer_module(A__)
												lowercase						     =						self.shortcut(A__)
												hidden_state += residual
												lowercase						     =						self.activation(A__)
												return hidden_state







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




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

												lowercase						     =						TFRegNetXLayer if config.layer_type == '''x''' else TFRegNetYLayer
												lowercase						     =						[
												    # downsampling is done in the first layer with stride of 2
												    layer(A__       ,A__       ,A__       ,stride=A__       ,name='''layers.0'''),
												    *[layer(A__       ,A__       ,A__       ,name=f'layers.{i+1}') for i in range(depth - 1)],
												]





							def 			A__				(    self       ,A__):
												for layer_module in self.layers:
																	lowercase						     =						layer_module(A__)
												return hidden_state







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




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





							def 			A__				(    self       ,A__       ,A__ = False       ,A__ = True):
												lowercase						     =						() if output_hidden_states else None

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

																	lowercase						     =						stage_module(A__)

												if output_hidden_states:
																	lowercase						     =						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=A__       ,hidden_states=A__)







@keras_serializable
class     lowercase		( tf.keras.layers.Layer		):
							lowercase_	:	int	 =RegNetConfig




							def __init__(    self       ,A__       ,**A__):
												super().__init__(**A__)
												lowercase						     =						config
												lowercase						     =						TFRegNetEmbeddings(A__       ,name='''embedder''')
												lowercase						     =						TFRegNetEncoder(A__       ,name='''encoder''')
												lowercase						     =						tf.keras.layers.GlobalAveragePoolingaD(keepdims=A__       ,name='''pooler''')





							@unpack_inputs
							def 			A__				(    self       ,A__       ,A__ = None       ,A__ = None       ,A__ = False       ,):
												lowercase						     =						(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												lowercase						     =						return_dict if return_dict is not None else self.config.use_return_dict

												lowercase						     =						self.embedder(A__       ,training=A__)

												lowercase						     =						self.encoder(
												    A__       ,output_hidden_states=A__       ,return_dict=A__       ,training=A__)

												lowercase						     =						encoder_outputs[0]
												lowercase						     =						self.pooler(A__)

												# Change to NCHW output format have uniformity in the modules
												lowercase						     =						tf.transpose(A__       ,perm=(0, 3, 1, 2))
												lowercase						     =						tf.transpose(A__       ,perm=(0, 3, 1, 2))

												# Change the other hidden state outputs to NCHW as well
												if output_hidden_states:
																	lowercase						     =						tuple([tf.transpose(A__       ,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=A__       ,pooler_output=A__       ,hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states       ,)







class     lowercase		( SCREAMING_SNAKE_CASE__		):
							lowercase_	:	List[str]	 =RegNetConfig
							lowercase_	:	Any	 ='''regnet'''
							lowercase_	:	str	 ='''pixel_values'''




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


lowercase__		:List[str]			  =   r"\n    Parameters:\n    This model is a Tensorflow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and\n    behavior.\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n"

lowercase__		:List[Any]			  =   r"\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConveNextImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"







@add_start_docstrings(
    '''The bare RegNet model outputting raw features without any specific head on top.'''	,  SCREAMING_SNAKE_CASE__	,  )
class     lowercase		( SCREAMING_SNAKE_CASE__		):




							def __init__(    self       ,A__       ,*A__       ,**A__):
												super().__init__(A__       ,*A__       ,**A__)
												lowercase						     =						TFRegNetMainLayer(A__       ,name='''regnet''')





							@unpack_inputs
							@add_start_docstrings_to_model_forward(A__)
							@add_code_sample_docstrings(
							    checkpoint=_CHECKPOINT_FOR_DOC       ,output_type=A__       ,config_class=_CONFIG_FOR_DOC       ,modality='''vision'''       ,expected_output=_EXPECTED_OUTPUT_SHAPE       ,)
							def 			A__				(    self       ,A__       ,A__ = None       ,A__ = None       ,A__=False       ,):
												lowercase						     =						(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												lowercase						     =						return_dict if return_dict is not None else self.config.use_return_dict

												lowercase						     =						self.regnet(
												    pixel_values=A__       ,output_hidden_states=A__       ,return_dict=A__       ,training=A__       ,)
												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.
    '''	,  SCREAMING_SNAKE_CASE__	,  )
class     lowercase		( SCREAMING_SNAKE_CASE__	,  SCREAMING_SNAKE_CASE__		):




							def __init__(    self       ,A__       ,*A__       ,**A__):
												super().__init__(A__       ,*A__       ,**A__)
												lowercase						     =						config.num_labels
												lowercase						     =						TFRegNetMainLayer(A__       ,name='''regnet''')
												# classification head
												lowercase						     =						[
												    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(A__)
							@add_code_sample_docstrings(
							    checkpoint=_IMAGE_CLASS_CHECKPOINT       ,output_type=A__       ,config_class=_CONFIG_FOR_DOC       ,expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT       ,)
							def 			A__				(    self       ,A__ = None       ,A__ = None       ,A__ = None       ,A__ = None       ,A__=False       ,):
												lowercase						     =						(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												lowercase						     =						return_dict if return_dict is not None else self.config.use_return_dict

												lowercase						     =						self.regnet(
												    A__       ,output_hidden_states=A__       ,return_dict=A__       ,training=A__)

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

												lowercase						     =						self.classifier[0](A__)
												lowercase						     =						self.classifier[1](A__)

												lowercase						     =						None if labels is None else self.hf_compute_loss(labels=A__       ,logits=A__)

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

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






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


lowercase__		:str			  =   {"configuration_vit": ["VIT_PRETRAINED_CONFIG_ARCHIVE_MAP", "ViTConfig", "ViTOnnxConfig"]}

try:
							if not is_vision_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowercase__		:str			  =   ["ViTFeatureExtractor"]
							lowercase__		:int			  =   ["ViTImageProcessor"]

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowercase__		:int			  =   [
							    "VIT_PRETRAINED_MODEL_ARCHIVE_LIST",
							    "ViTForImageClassification",
							    "ViTForMaskedImageModeling",
							    "ViTModel",
							    "ViTPreTrainedModel",
							]

try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowercase__		:Any			  =   [
							    "TFViTForImageClassification",
							    "TFViTModel",
							    "TFViTPreTrainedModel",
							]

try:
							if not is_flax_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowercase__		:Optional[int]			  =   [
							    "FlaxViTForImageClassification",
							    "FlaxViTModel",
							    "FlaxViTPreTrainedModel",
							]

if TYPE_CHECKING:
							from .configuration_vit import VIT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTConfig, ViTOnnxConfig

							try:
														if not is_vision_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .feature_extraction_vit import ViTFeatureExtractor
														from .image_processing_vit import ViTImageProcessor

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_vit import (
														    VIT_PRETRAINED_MODEL_ARCHIVE_LIST,
														    ViTForImageClassification,
														    ViTForMaskedImageModeling,
														    ViTModel,
														    ViTPreTrainedModel,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tf_vit import TFViTForImageClassification, TFViTModel, TFViTPreTrainedModel

							try:
														if not is_flax_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel, FlaxViTPreTrainedModel


else:
							import sys

							lowercase__		:Dict			  =   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




