




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


UpperCamelCase								=	logging.get_logger(__name__)

# General docstring
UpperCamelCase								=	'RegNetConfig'

# Base docstring
UpperCamelCase								=	'facebook/regnet-y-040'
UpperCamelCase								=	[1, 1088, 7, 7]

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

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





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




    """simple docstring"""

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


    def   a (			self	:	Optional[int]  ,      SCREAMING_SNAKE_CASE__	:	str  ) ->			Dict:
         lowerCAmelCase__            =					self.convolution(self.padding(__snake_case  )  )
         lowerCAmelCase__            =					self.normalization(__snake_case  )
         lowerCAmelCase__            =					self.activation(__snake_case  )
         return hidden_state






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




    """simple docstring"""

    def __init__(			self	:	Union[str, Any]  ,      SCREAMING_SNAKE_CASE__	:	RegNetConfig  ,      **SCREAMING_SNAKE_CASE__	:	Dict  ) ->			Optional[int]:
         super().__init__(**__snake_case  )
         lowerCAmelCase__            =					config.num_channels
         lowerCAmelCase__            =					TFRegNetConvLayer(
             out_channels=config.embedding_size  ,      kernel_size=3  ,      stride=2  ,      activation=config.hidden_act  ,      name="embedder"  ,      )


    def   a (			self	:	List[str]  ,      SCREAMING_SNAKE_CASE__	:	Dict  ) ->			int:
         lowerCAmelCase__            =					shape_list(__snake_case  )[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__            =					tf.transpose(__snake_case  ,      perm=(0, 2, 3, 1)  )
         lowerCAmelCase__            =					self.embedder(__snake_case  )
         return hidden_state






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




    """simple docstring"""

    def __init__(			self	:	Union[str, Any]  ,      SCREAMING_SNAKE_CASE__	:	int  ,      SCREAMING_SNAKE_CASE__	:	int = 2  ,      **SCREAMING_SNAKE_CASE__	:	Any  ) ->			Dict:
         super().__init__(**__snake_case  )
         lowerCAmelCase__            =					tf.keras.layers.ConvaD(
             filters=__snake_case  ,      kernel_size=1  ,      strides=__snake_case  ,      use_bias=__snake_case  ,      name="convolution"  )
         lowerCAmelCase__            =					tf.keras.layers.BatchNormalization(epsilon=1e-5  ,      momentum=0.9  ,      name="normalization"  )


    def   a (			self	:	Optional[Any]  ,      SCREAMING_SNAKE_CASE__	:	tf.Tensor  ,      SCREAMING_SNAKE_CASE__	:	bool = False  ) ->			tf.Tensor:
         return self.normalization(self.convolution(__snake_case  )  ,      training=__snake_case  )






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




    """simple docstring"""

    def __init__(			self	:	int  ,      SCREAMING_SNAKE_CASE__	:	int  ,      SCREAMING_SNAKE_CASE__	:	int  ,      **SCREAMING_SNAKE_CASE__	:	str  ) ->			str:
         super().__init__(**__snake_case  )
         lowerCAmelCase__            =					tf.keras.layers.GlobalAveragePoolingaD(keepdims=__snake_case  ,      name="pooler"  )
         lowerCAmelCase__            =					[
             tf.keras.layers.ConvaD(filters=__snake_case  ,      kernel_size=1  ,      activation="relu"  ,      name="attention.0"  ),
             tf.keras.layers.ConvaD(filters=__snake_case  ,      kernel_size=1  ,      activation="sigmoid"  ,      name="attention.2"  ),
         ]


    def   a (			self	:	Dict  ,      SCREAMING_SNAKE_CASE__	:	List[str]  ) ->			List[Any]:
         # [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
         lowerCAmelCase__            =					self.pooler(__snake_case  )
         for layer_module in self.attention:
              lowerCAmelCase__            =					layer_module(__snake_case  )
         lowerCAmelCase__            =					hidden_state * pooled
         return hidden_state






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




    """simple docstring"""

    def __init__(			self	:	Union[str, Any]  ,      SCREAMING_SNAKE_CASE__	:	RegNetConfig  ,      SCREAMING_SNAKE_CASE__	:	int  ,      SCREAMING_SNAKE_CASE__	:	int  ,      SCREAMING_SNAKE_CASE__	:	int = 1  ,      **SCREAMING_SNAKE_CASE__	:	Optional[int]  ) ->			Optional[int]:
         super().__init__(**__snake_case  )
         lowerCAmelCase__            =					in_channels != out_channels or stride != 1
         lowerCAmelCase__            =					max(1  ,      out_channels // config.groups_width  )
         lowerCAmelCase__            =					(
             TFRegNetShortCut(__snake_case  ,      stride=__snake_case  ,      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__            =					[
             TFRegNetConvLayer(__snake_case  ,      kernel_size=1  ,      activation=config.hidden_act  ,      name="layer.0"  ),
             TFRegNetConvLayer(
                 __snake_case  ,      stride=__snake_case  ,      groups=__snake_case  ,      activation=config.hidden_act  ,      name="layer.1"  ),
             TFRegNetConvLayer(__snake_case  ,      kernel_size=1  ,      activation=__snake_case  ,      name="layer.2"  ),
         ]
         lowerCAmelCase__            =					ACTaFN[config.hidden_act]


    def   a (			self	:	Optional[int]  ,      SCREAMING_SNAKE_CASE__	:	Any  ) ->			Union[str, Any]:
         lowerCAmelCase__            =					hidden_state
         for layer_module in self.layers:
              lowerCAmelCase__            =					layer_module(__snake_case  )
         lowerCAmelCase__            =					self.shortcut(__snake_case  )
         hidden_state += residual
         lowerCAmelCase__            =					self.activation(__snake_case  )
         return hidden_state






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




    """simple docstring"""

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


    def   a (			self	:	List[str]  ,      SCREAMING_SNAKE_CASE__	:	int  ) ->			Dict:
         lowerCAmelCase__            =					hidden_state
         for layer_module in self.layers:
              lowerCAmelCase__            =					layer_module(__snake_case  )
         lowerCAmelCase__            =					self.shortcut(__snake_case  )
         hidden_state += residual
         lowerCAmelCase__            =					self.activation(__snake_case  )
         return hidden_state






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




    """simple docstring"""

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

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


    def   a (			self	:	int  ,      SCREAMING_SNAKE_CASE__	:	Union[str, Any]  ) ->			Tuple:
         for layer_module in self.layers:
              lowerCAmelCase__            =					layer_module(__snake_case  )
         return hidden_state






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




    """simple docstring"""

    def __init__(			self	:	Union[str, Any]  ,      SCREAMING_SNAKE_CASE__	:	RegNetConfig  ,      **SCREAMING_SNAKE_CASE__	:	Optional[Any]  ) ->			Dict:
         super().__init__(**__snake_case  )
         lowerCAmelCase__            =					[]
         # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
         self.stages.append(
             TFRegNetStage(
                 __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"  ,      )  )
         lowerCAmelCase__            =					zip(config.hidden_sizes  ,      config.hidden_sizes[1:]  )
         for i, ((in_channels, out_channels), depth) in enumerate(zip(__snake_case  ,      config.depths[1:]  )  ):
              self.stages.append(TFRegNetStage(__snake_case  ,      __snake_case  ,      __snake_case  ,      depth=__snake_case  ,      name=f'stages.{i+1}'  )  )


    def   a (			self	:	int  ,      SCREAMING_SNAKE_CASE__	:	tf.Tensor  ,      SCREAMING_SNAKE_CASE__	:	bool = False  ,      SCREAMING_SNAKE_CASE__	:	bool = True  ) ->			TFBaseModelOutputWithNoAttention:
         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  )

         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  )






@keras_serializable
class 					__lowerCamelCase			(       tf.keras.layers.Layer       ):




    """simple docstring"""

    snake_case__   = RegNetConfig


    def __init__(			self	:	Optional[int]  ,      SCREAMING_SNAKE_CASE__	:	Any  ,      **SCREAMING_SNAKE_CASE__	:	List[str]  ) ->			int:
         super().__init__(**__snake_case  )
         lowerCAmelCase__            =					config
         lowerCAmelCase__            =					TFRegNetEmbeddings(__snake_case  ,      name="embedder"  )
         lowerCAmelCase__            =					TFRegNetEncoder(__snake_case  ,      name="encoder"  )
         lowerCAmelCase__            =					tf.keras.layers.GlobalAveragePoolingaD(keepdims=__snake_case  ,      name="pooler"  )


    @unpack_inputs
    def   a (			self	:	Optional[Any]  ,      SCREAMING_SNAKE_CASE__	:	tf.Tensor  ,      SCREAMING_SNAKE_CASE__	:	Optional[bool] = None  ,      SCREAMING_SNAKE_CASE__	:	Optional[bool] = None  ,      SCREAMING_SNAKE_CASE__	:	bool = False  ,      ) ->			TFBaseModelOutputWithPoolingAndNoAttention:
         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.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  )

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

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






class 					__lowerCamelCase			(       SCREAMING_SNAKE_CASE_       ):




    """simple docstring"""

    snake_case__   = RegNetConfig
    snake_case__   = "regnet"
    snake_case__   = "pixel_values"


    @property
    def   a (			self	:	Optional[int]  ) ->			Union[str, Any]:
         return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)  ,      dtype=tf.floataa  )}






UpperCamelCase								=	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'

UpperCamelCase								=	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 					__lowerCamelCase			(       SCREAMING_SNAKE_CASE_       ):




    """simple docstring"""

    def __init__(			self	:	Optional[Any]  ,      SCREAMING_SNAKE_CASE__	:	RegNetConfig  ,      *SCREAMING_SNAKE_CASE__	:	List[Any]  ,      **SCREAMING_SNAKE_CASE__	:	Tuple  ) ->			Tuple:
         super().__init__(__snake_case  ,      *__snake_case  ,      **__snake_case  )
         lowerCAmelCase__            =					TFRegNetMainLayer(__snake_case  ,      name="regnet"  )


    @unpack_inputs
    @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  ,      )
    def   a (			self	:	Dict  ,      SCREAMING_SNAKE_CASE__	:	tf.Tensor  ,      SCREAMING_SNAKE_CASE__	:	Optional[bool] = None  ,      SCREAMING_SNAKE_CASE__	:	Optional[bool] = None  ,      SCREAMING_SNAKE_CASE__	:	int=False  ,      ) ->			Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:
         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.regnet(
             pixel_values=__snake_case  ,      output_hidden_states=__snake_case  ,      return_dict=__snake_case  ,      training=__snake_case  ,      )
         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(
    "\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "      ,							SCREAMING_SNAKE_CASE_      ,							)
class 					__lowerCamelCase			(       SCREAMING_SNAKE_CASE_      ,							SCREAMING_SNAKE_CASE_       ):




    """simple docstring"""

    def __init__(			self	:	int  ,      SCREAMING_SNAKE_CASE__	:	RegNetConfig  ,      *SCREAMING_SNAKE_CASE__	:	Any  ,      **SCREAMING_SNAKE_CASE__	:	Any  ) ->			Optional[Any]:
         super().__init__(__snake_case  ,      *__snake_case  ,      **__snake_case  )
         lowerCAmelCase__            =					config.num_labels
         lowerCAmelCase__            =					TFRegNetMainLayer(__snake_case  ,      name="regnet"  )
         # classification head
         lowerCAmelCase__            =					[
             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(__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  ,      )
    def   a (			self	:	List[str]  ,      SCREAMING_SNAKE_CASE__	:	tf.Tensor = None  ,      SCREAMING_SNAKE_CASE__	:	tf.Tensor = None  ,      SCREAMING_SNAKE_CASE__	:	bool = None  ,      SCREAMING_SNAKE_CASE__	:	bool = None  ,      SCREAMING_SNAKE_CASE__	:	Dict=False  ,      ) ->			Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
         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.regnet(
             __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[0](__snake_case  )
         lowerCAmelCase__            =					self.classifier[1](__snake_case  )

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

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

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








"""simple docstring"""





def   a    (     __UpperCAmelCase      :    list[int]  )    ->							float:
						if not nums:  # Makes sure that the list is not empty
												raise ValueError("""List is empty"""  )

						__magic_name__:	Dict							=       sum(__UpperCAmelCase  ) / len(__UpperCAmelCase  )  # Calculate the average
						return sum(abs(x - average  ) for x in nums  ) / len(__UpperCAmelCase  )


if __name__ == "__main__":
							import doctest

							doctest.testmod()



