



"""simple docstring"""
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						:		Tuple			= logging.get_logger(__name__)

# General docstring
UpperCAmelCase						:		Union[str, Any]			= "RegNetConfig"

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

# Image classification docstring
UpperCAmelCase						:		Optional[Any]			= "facebook/regnet-y-040"
UpperCAmelCase						:		Any			= "tabby, tabby cat"

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


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

    def __init__(				self			:				Union[str, Any] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int = 3 ,  lowerCAmelCase_			:				int = 1 ,  lowerCAmelCase_			:				int = 1 ,  lowerCAmelCase_			:				Optional[str] = "relu" ,  **lowerCAmelCase_			:				Dict ,  ):







      """simple docstring"""







      super().__init__(**lowerCAmelCase_)
      # 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=lowerCAmelCase_ ,  kernel_size=lowerCAmelCase_ ,  strides=lowerCAmelCase_ ,  padding="""VALID""" ,  groups=lowerCAmelCase_ ,  use_bias=lowerCAmelCase_ ,  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 					_UpperCAmelCase	(				self			:				Optional[int] ,  lowerCAmelCase_			:				int):







      """simple docstring"""







      lowercase_								= self.convolution(self.padding(lowerCAmelCase_))
      lowercase_								= self.normalization(lowerCAmelCase_)
      lowercase_								= self.activation(lowerCAmelCase_)
      return hidden_state







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

    def __init__(				self			:				List[Any] ,  lowerCAmelCase_			:				RegNetConfig ,  **lowerCAmelCase_			:				Dict):







      """simple docstring"""







      super().__init__(**lowerCAmelCase_)
      lowercase_								= config.num_channels
      lowercase_								= TFRegNetConvLayer(
          out_channels=config.embedding_size ,  kernel_size=3 ,  stride=2 ,  activation=config.hidden_act ,  name="""embedder""" ,  )




    def 					_UpperCAmelCase	(				self			:				Any ,  lowerCAmelCase_			:				Tuple):







      """simple docstring"""







      lowercase_								= shape_list(lowerCAmelCase_)[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(lowerCAmelCase_ ,  perm=(0, 2, 3, 1))
      lowercase_								= self.embedder(lowerCAmelCase_)
      return hidden_state







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

    def __init__(				self			:				Tuple ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int = 2 ,  **lowerCAmelCase_			:				List[str]):







      """simple docstring"""







      super().__init__(**lowerCAmelCase_)
      lowercase_								= tf.keras.layers.ConvaD(
          filters=lowerCAmelCase_ ,  kernel_size=1 ,  strides=lowerCAmelCase_ ,  use_bias=lowerCAmelCase_ ,  name="""convolution""")
      lowercase_								= tf.keras.layers.BatchNormalization(epsilon=1E-5 ,  momentum=0.9 ,  name="""normalization""")




    def 					_UpperCAmelCase	(				self			:				Dict ,  lowerCAmelCase_			:				tf.Tensor ,  lowerCAmelCase_			:				bool = False):







      """simple docstring"""







      return self.normalization(self.convolution(lowerCAmelCase_) ,  training=lowerCAmelCase_)







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

    def __init__(				self			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int ,  **lowerCAmelCase_			:				Dict):







      """simple docstring"""







      super().__init__(**lowerCAmelCase_)
      lowercase_								= tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowerCAmelCase_ ,  name="""pooler""")
      lowercase_								= [
          tf.keras.layers.ConvaD(filters=lowerCAmelCase_ ,  kernel_size=1 ,  activation="""relu""" ,  name="""attention.0"""),
          tf.keras.layers.ConvaD(filters=lowerCAmelCase_ ,  kernel_size=1 ,  activation="""sigmoid""" ,  name="""attention.2"""),
      ]




    def 					_UpperCAmelCase	(				self			:				Tuple ,  lowerCAmelCase_			:				Any):







      """simple docstring"""







      lowercase_								= self.pooler(lowerCAmelCase_)
      for layer_module in self.attention:
        lowercase_								= layer_module(lowerCAmelCase_)
      lowercase_								= hidden_state * pooled
      return hidden_state







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

    def __init__(				self			:				Tuple ,  lowerCAmelCase_			:				RegNetConfig ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int = 1 ,  **lowerCAmelCase_			:				Dict):







      """simple docstring"""







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




    def 					_UpperCAmelCase	(				self			:				Union[str, Any] ,  lowerCAmelCase_			:				Optional[int]):







      """simple docstring"""







      lowercase_								= hidden_state
      for layer_module in self.layers:
        lowercase_								= layer_module(lowerCAmelCase_)
      lowercase_								= self.shortcut(lowerCAmelCase_)
      hidden_state += residual
      lowercase_								= self.activation(lowerCAmelCase_)
      return hidden_state







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

    def __init__(				self			:				Dict ,  lowerCAmelCase_			:				RegNetConfig ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int = 1 ,  **lowerCAmelCase_			:				Optional[Any]):







      """simple docstring"""







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




    def 					_UpperCAmelCase	(				self			:				Tuple ,  lowerCAmelCase_			:				List[str]):







      """simple docstring"""







      lowercase_								= hidden_state
      for layer_module in self.layers:
        lowercase_								= layer_module(lowerCAmelCase_)
      lowercase_								= self.shortcut(lowerCAmelCase_)
      hidden_state += residual
      lowercase_								= self.activation(lowerCAmelCase_)
      return hidden_state







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

    def __init__(				self			:				Tuple ,  lowerCAmelCase_			:				RegNetConfig ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int = 2 ,  lowerCAmelCase_			:				int = 2 ,  **lowerCAmelCase_			:				int):







      """simple docstring"""







      super().__init__(**lowerCAmelCase_)

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




    def 					_UpperCAmelCase	(				self			:				List[str] ,  lowerCAmelCase_			:				int):







      """simple docstring"""







      for layer_module in self.layers:
        lowercase_								= layer_module(lowerCAmelCase_)
      return hidden_state







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

    def __init__(				self			:				List[str] ,  lowerCAmelCase_			:				RegNetConfig ,  **lowerCAmelCase_			:				Optional[int]):







      """simple docstring"""







      super().__init__(**lowerCAmelCase_)
      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(
              lowerCAmelCase_ ,  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(lowerCAmelCase_ ,  config.depths[1:])):
        self.stages.append(TFRegNetStage(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  depth=lowerCAmelCase_ ,  name=F'''stages.{i+1}'''))




    def 					_UpperCAmelCase	(				self			:				Dict ,  lowerCAmelCase_			:				tf.Tensor ,  lowerCAmelCase_			:				bool = False ,  lowerCAmelCase_			:				bool = True):







      """simple docstring"""







      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(lowerCAmelCase_)

      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=lowerCAmelCase_ ,  hidden_states=lowerCAmelCase_)







@keras_serializable
class    SCREAMING_SNAKE_CASE__ (							tf.keras.layers.Layer   ):
    lowercase__      =   RegNetConfig

    def __init__(				self			:				Any ,  lowerCAmelCase_			:				Any ,  **lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      super().__init__(**lowerCAmelCase_)
      lowercase_								= config
      lowercase_								= TFRegNetEmbeddings(lowerCAmelCase_ ,  name="""embedder""")
      lowercase_								= TFRegNetEncoder(lowerCAmelCase_ ,  name="""encoder""")
      lowercase_								= tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowerCAmelCase_ ,  name="""pooler""")




    @unpack_inputs
    def 					_UpperCAmelCase	(				self			:				Any ,  lowerCAmelCase_			:				tf.Tensor ,  lowerCAmelCase_			:				Optional[bool] = None ,  lowerCAmelCase_			:				Optional[bool] = None ,  lowerCAmelCase_			:				bool = False ,  ):







      """simple docstring"""







      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(lowerCAmelCase_ ,  training=lowerCAmelCase_)

      lowercase_								= self.encoder(
          lowerCAmelCase_ ,  output_hidden_states=lowerCAmelCase_ ,  return_dict=lowerCAmelCase_ ,  training=lowerCAmelCase_)

      lowercase_								= encoder_outputs[0]
      lowercase_								= self.pooler(lowerCAmelCase_)

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

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







class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase   ):
    lowercase__      =   RegNetConfig
    lowercase__      =   "regnet"
    lowercase__      =   "pixel_values"

    @property
    def 					_UpperCAmelCase	(				self			:				List[Any]):







      """simple docstring"""







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







UpperCAmelCase						:		Optional[int]			= 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						:		Tuple			= 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."      ,  __UpperCAmelCase      ,  )
class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase   ):

    def __init__(				self			:				List[Any] ,  lowerCAmelCase_			:				RegNetConfig ,  *lowerCAmelCase_			:				Any ,  **lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      super().__init__(lowerCAmelCase_ ,  *lowerCAmelCase_ ,  **lowerCAmelCase_)
      lowercase_								= TFRegNetMainLayer(lowerCAmelCase_ ,  name="""regnet""")




    @unpack_inputs
    @add_start_docstrings_to_model_forward(lowerCAmelCase_)
    @add_code_sample_docstrings(
        checkpoint=_CHECKPOINT_FOR_DOC ,  output_type=lowerCAmelCase_ ,  config_class=_CONFIG_FOR_DOC ,  modality="""vision""" ,  expected_output=_EXPECTED_OUTPUT_SHAPE ,  )
    def 					_UpperCAmelCase	(				self			:				Optional[int] ,  lowerCAmelCase_			:				tf.Tensor ,  lowerCAmelCase_			:				Optional[bool] = None ,  lowerCAmelCase_			:				Optional[bool] = None ,  lowerCAmelCase_			:				int=False ,  ):







      """simple docstring"""







      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=lowerCAmelCase_ ,  output_hidden_states=lowerCAmelCase_ ,  return_dict=lowerCAmelCase_ ,  training=lowerCAmelCase_ ,  )
      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    "      ,  __UpperCAmelCase      ,  )
class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  __UpperCAmelCase   ):

    def __init__(				self			:				Dict ,  lowerCAmelCase_			:				RegNetConfig ,  *lowerCAmelCase_			:				Tuple ,  **lowerCAmelCase_			:				str):







      """simple docstring"""







      super().__init__(lowerCAmelCase_ ,  *lowerCAmelCase_ ,  **lowerCAmelCase_)
      lowercase_								= config.num_labels
      lowercase_								= TFRegNetMainLayer(lowerCAmelCase_ ,  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(lowerCAmelCase_)
    @add_code_sample_docstrings(
        checkpoint=_IMAGE_CLASS_CHECKPOINT ,  output_type=lowerCAmelCase_ ,  config_class=_CONFIG_FOR_DOC ,  expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT ,  )
    def 					_UpperCAmelCase	(				self			:				int ,  lowerCAmelCase_			:				tf.Tensor = None ,  lowerCAmelCase_			:				tf.Tensor = None ,  lowerCAmelCase_			:				bool = None ,  lowerCAmelCase_			:				bool = None ,  lowerCAmelCase_			:				Tuple=False ,  ):







      """simple docstring"""







      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(
          lowerCAmelCase_ ,  output_hidden_states=lowerCAmelCase_ ,  return_dict=lowerCAmelCase_ ,  training=lowerCAmelCase_)

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

      lowercase_								= self.classifier[0](lowerCAmelCase_)
      lowercase_								= self.classifier[1](lowerCAmelCase_)

      lowercase_								= None if labels is None else self.hf_compute_loss(labels=lowerCAmelCase_ ,  logits=lowerCAmelCase_)

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

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






"""simple docstring"""
from ..utils import DummyObject, requires_backends


class    SCREAMING_SNAKE_CASE__ (							metaclass=__UpperCAmelCase   ):
    lowercase__      =   ["keras_nlp"]

    def __init__(				self			:				Dict ,  *lowerCAmelCase_			:				List[str] ,  **lowerCAmelCase_			:				Tuple):







      """simple docstring"""







      requires_backends(self ,  ["""keras_nlp"""])

