





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


snake_case_     =	logging.get_logger(__name__)

# General docstring
snake_case_     =	"""RegNetConfig"""

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

# Image classification docstring
snake_case_     =	"""facebook/regnet-y-040"""
snake_case_     =	"""tabby, tabby cat"""

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


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


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


   def 					snake_case				( self    ,   _UpperCAmelCase					):
       A       :       Tuple							=						self.convolution(self.padding(_UpperCAmelCase					)					)
       A       :       List[Any]							=						self.normalization(_UpperCAmelCase					)
       A       :       List[Any]							=						self.activation(_UpperCAmelCase					)
       return hidden_state







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


   def __init__( self    ,   _UpperCAmelCase    ,   **_UpperCAmelCase					):
       super().__init__(**_UpperCAmelCase					)
       A       :       int							=						config.num_channels
       A       :       List[Any]							=						TFRegNetConvLayer(
           out_channels=config.embedding_size    ,   kernel_size=3    ,   stride=2    ,   activation=config.hidden_act    ,   name='''embedder'''    ,   )


   def 					snake_case				( self    ,   _UpperCAmelCase					):
       A       :       Union[str, Any]							=						shape_list(_UpperCAmelCase					)[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)
       A       :       Optional[Any]							=						tf.transpose(_UpperCAmelCase    ,   perm=(0, 2, 3, 1)					)
       A       :       Dict							=						self.embedder(_UpperCAmelCase					)
       return hidden_state







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


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


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







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


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


   def 					snake_case				( self    ,   _UpperCAmelCase					):
       # [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
       A       :       Dict							=						self.pooler(_UpperCAmelCase					)
       for layer_module in self.attention:
           A       :       List[str]							=						layer_module(_UpperCAmelCase					)
       A       :       Dict							=						hidden_state * pooled
       return hidden_state







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


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


   def 					snake_case				( self    ,   _UpperCAmelCase					):
       A       :       int							=						hidden_state
       for layer_module in self.layers:
           A       :       Any							=						layer_module(_UpperCAmelCase					)
       A       :       Optional[Any]							=						self.shortcut(_UpperCAmelCase					)
       hidden_state += residual
       A       :       List[Any]							=						self.activation(_UpperCAmelCase					)
       return hidden_state







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


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


   def 					snake_case				( self    ,   _UpperCAmelCase					):
       A       :       Union[str, Any]							=						hidden_state
       for layer_module in self.layers:
           A       :       str							=						layer_module(_UpperCAmelCase					)
       A       :       Optional[Any]							=						self.shortcut(_UpperCAmelCase					)
       hidden_state += residual
       A       :       str							=						self.activation(_UpperCAmelCase					)
       return hidden_state







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


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

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


   def 					snake_case				( self    ,   _UpperCAmelCase					):
       for layer_module in self.layers:
           A       :       int							=						layer_module(_UpperCAmelCase					)
       return hidden_state







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


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


   def 					snake_case				( self    ,   _UpperCAmelCase    ,   _UpperCAmelCase = False    ,   _UpperCAmelCase = True					):
       A       :       Any							=						() if output_hidden_states else None

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

           A       :       List[str]							=						stage_module(_UpperCAmelCase					)

       if output_hidden_states:
           A       :       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=_UpperCAmelCase    ,   hidden_states=_UpperCAmelCase					)







@keras_serializable
class   _lowercase						(							tf.keras.layers.Layer				):
   _UpperCamelCase       		=   RegNetConfig


   def __init__( self    ,   _UpperCAmelCase    ,   **_UpperCAmelCase					):
       super().__init__(**_UpperCAmelCase					)
       A       :       Optional[int]							=						config
       A       :       Dict							=						TFRegNetEmbeddings(_UpperCAmelCase    ,   name='''embedder'''					)
       A       :       List[str]							=						TFRegNetEncoder(_UpperCAmelCase    ,   name='''encoder'''					)
       A       :       List[Any]							=						tf.keras.layers.GlobalAveragePoolingaD(keepdims=_UpperCAmelCase    ,   name='''pooler'''					)


   @unpack_inputs
   def 					snake_case				( self    ,   _UpperCAmelCase    ,   _UpperCAmelCase = None    ,   _UpperCAmelCase = None    ,   _UpperCAmelCase = False    ,   ):
       A       :       int							=						(
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )
       A       :       str							=						return_dict if return_dict is not None else self.config.use_return_dict

       A       :       List[Any]							=						self.embedder(_UpperCAmelCase    ,   training=_UpperCAmelCase					)

       A       :       Optional[Any]							=						self.encoder(
           _UpperCAmelCase    ,   output_hidden_states=_UpperCAmelCase    ,   return_dict=_UpperCAmelCase    ,   training=_UpperCAmelCase					)

       A       :       List[str]							=						encoder_outputs[0]
       A       :       Union[str, Any]							=						self.pooler(_UpperCAmelCase					)

       # Change to NCHW output format have uniformity in the modules
       A       :       int							=						tf.transpose(_UpperCAmelCase    ,   perm=(0, 3, 1, 2)					)
       A       :       Optional[int]							=						tf.transpose(_UpperCAmelCase    ,   perm=(0, 3, 1, 2)					)

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







class   _lowercase						(							a				):
   _UpperCamelCase       		=   RegNetConfig
   _UpperCamelCase       		=   """regnet"""
   _UpperCamelCase       		=   """pixel_values"""


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







snake_case_     =	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.
"""

snake_case_     =	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   _lowercase						(							a				):


   def __init__( self    ,   _UpperCAmelCase    ,   *_UpperCAmelCase    ,   **_UpperCAmelCase					):
       super().__init__(_UpperCAmelCase    ,   *_UpperCAmelCase    ,   **_UpperCAmelCase					)
       A       :       int							=						TFRegNetMainLayer(_UpperCAmelCase    ,   name='''regnet'''					)


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

       A       :       Optional[Any]							=						self.regnet(
           pixel_values=_UpperCAmelCase    ,   output_hidden_states=_UpperCAmelCase    ,   return_dict=_UpperCAmelCase    ,   training=_UpperCAmelCase    ,   )
       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   _lowercase						(							a		,					a				):


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

       A       :       Tuple							=						self.regnet(
           _UpperCAmelCase    ,   output_hidden_states=_UpperCAmelCase    ,   return_dict=_UpperCAmelCase    ,   training=_UpperCAmelCase					)

       A       :       Any							=						outputs.pooler_output if return_dict else outputs[1]

       A       :       int							=						self.classifier[0](_UpperCAmelCase					)
       A       :       Any							=						self.classifier[1](_UpperCAmelCase					)

       A       :       Dict							=						None if labels is None else self.hf_compute_loss(labels=_UpperCAmelCase    ,   logits=_UpperCAmelCase					)

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

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







'''simple docstring'''






from ..utils import DummyObject, requires_backends


class   _lowercase						(							metaclass=a				):
   _UpperCamelCase       		=   ["""onnx"""]


   def __init__( self    ,   *_UpperCAmelCase    ,   **_UpperCAmelCase					):
       requires_backends(self    ,   ['''onnx''']					)


   @classmethod
   def 					snake_case				( cls    ,   *_UpperCAmelCase    ,   **_UpperCAmelCase					):
       requires_backends(cls    ,   ['''onnx''']					)


   @classmethod
   def 					snake_case				( cls    ,   *_UpperCAmelCase    ,   **_UpperCAmelCase					):
       requires_backends(cls    ,   ['''onnx''']					)
