


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


lowercase							=logging.get_logger(__name__)

# General docstring
lowercase							='RegNetConfig'

# Base docstring
lowercase							='facebook/regnet-y-040'
lowercase							=[1, 1088, 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  __magic_name__							(		tf.keras.layers.Layer    ):






       def __init__(      self  ,							snake_case  ,							snake_case = 3  ,							snake_case = 1  ,							snake_case = 1  ,							snake_case = "relu"  ,							**snake_case  ,							)						->  List[str]:




          '''simple docstring'''
          super().__init__(**snake_case)
          # The padding and conv has been verified in
          # https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
          _UpperCAmelCase       :       str							   =tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2)
          _UpperCAmelCase       :       Dict							   =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'  ,							)
          _UpperCAmelCase       :       Dict							   =tf.keras.layers.BatchNormalization(epsilon=1E-5  ,							momentum=0.9  ,							name='normalization')
          _UpperCAmelCase       :       Tuple							   =ACTaFN[activation] if activation is not None else tf.identity






       def       lowerCAmelCase		(      self  ,							snake_case)						->  Dict:




          '''simple docstring'''
          _UpperCAmelCase       :       Union[str, Any]							   =self.convolution(self.padding(snake_case))
          _UpperCAmelCase       :       Dict							   =self.normalization(snake_case)
          _UpperCAmelCase       :       Dict							   =self.activation(snake_case)
          return hidden_state






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






       def __init__(      self  ,							snake_case  ,							**snake_case)						->  Any:




          '''simple docstring'''
          super().__init__(**snake_case)
          _UpperCAmelCase       :       int							   =config.num_channels
          _UpperCAmelCase       :       Optional[Any]							   =TFRegNetConvLayer(
              out_channels=config.embedding_size  ,							kernel_size=3  ,							stride=2  ,							activation=config.hidden_act  ,							name='embedder'  ,							)






       def       lowerCAmelCase		(      self  ,							snake_case)						->  Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase       :       List[str]							   =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)
          _UpperCAmelCase       :       Any							   =tf.transpose(snake_case  ,							perm=(0, 2, 3, 1))
          _UpperCAmelCase       :       int							   =self.embedder(snake_case)
          return hidden_state






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






       def __init__(      self  ,							snake_case  ,							snake_case = 2  ,							**snake_case)						->  int:




          '''simple docstring'''
          super().__init__(**snake_case)
          _UpperCAmelCase       :       Union[str, Any]							   =tf.keras.layers.ConvaD(
              filters=snake_case  ,							kernel_size=1  ,							strides=snake_case  ,							use_bias=snake_case  ,							name='convolution')
          _UpperCAmelCase       :       Optional[int]							   =tf.keras.layers.BatchNormalization(epsilon=1E-5  ,							momentum=0.9  ,							name='normalization')






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case = False)						->  tf.Tensor:




          '''simple docstring'''
          return self.normalization(self.convolution(snake_case)  ,							training=snake_case)






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






       def __init__(      self  ,							snake_case  ,							snake_case  ,							**snake_case)						->  Dict:




          '''simple docstring'''
          super().__init__(**snake_case)
          _UpperCAmelCase       :       Dict							   =tf.keras.layers.GlobalAveragePoolingaD(keepdims=snake_case  ,							name='pooler')
          _UpperCAmelCase       :       Union[str, Any]							   =[
              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       lowerCAmelCase		(      self  ,							snake_case)						->  Dict:




          '''simple docstring'''
          # [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
          _UpperCAmelCase       :       Union[str, Any]							   =self.pooler(snake_case)
          for layer_module in self.attention:
             _UpperCAmelCase       :       str							   =layer_module(snake_case)
          _UpperCAmelCase       :       Tuple							   =hidden_state * pooled
          return hidden_state






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






       def __init__(      self  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case = 1  ,							**snake_case)						->  Any:




          '''simple docstring'''
          super().__init__(**snake_case)
          _UpperCAmelCase       :       int							   =in_channels != out_channels or stride != 1
          _UpperCAmelCase       :       Union[str, Any]							   =max(1  ,							out_channels // config.groups_width)
          _UpperCAmelCase       :       Union[str, Any]							   =(
              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.
          _UpperCAmelCase       :       List[Any]							   =[
              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'),
          ]
          _UpperCAmelCase       :       List[str]							   =ACTaFN[config.hidden_act]






       def       lowerCAmelCase		(      self  ,							snake_case)						->  Any:




          '''simple docstring'''
          _UpperCAmelCase       :       Union[str, Any]							   =hidden_state
          for layer_module in self.layers:
             _UpperCAmelCase       :       int							   =layer_module(snake_case)
          _UpperCAmelCase       :       Dict							   =self.shortcut(snake_case)
          hidden_state += residual
          _UpperCAmelCase       :       Optional[Any]							   =self.activation(snake_case)
          return hidden_state






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






       def __init__(      self  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case = 1  ,							**snake_case)						->  Optional[Any]:




          '''simple docstring'''
          super().__init__(**snake_case)
          _UpperCAmelCase       :       Any							   =in_channels != out_channels or stride != 1
          _UpperCAmelCase       :       str							   =max(1  ,							out_channels // config.groups_width)
          _UpperCAmelCase       :       Dict							   =(
              TFRegNetShortCut(snake_case  ,							stride=snake_case  ,							name='shortcut')
              if should_apply_shortcut
              else tf.keras.layers.Activation('linear'  ,							name='shortcut')
          )
          _UpperCAmelCase       :       Optional[int]							   =[
              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'),
          ]
          _UpperCAmelCase       :       int							   =ACTaFN[config.hidden_act]






       def       lowerCAmelCase		(      self  ,							snake_case)						->  Tuple:




          '''simple docstring'''
          _UpperCAmelCase       :       Tuple							   =hidden_state
          for layer_module in self.layers:
             _UpperCAmelCase       :       Any							   =layer_module(snake_case)
          _UpperCAmelCase       :       List[Any]							   =self.shortcut(snake_case)
          hidden_state += residual
          _UpperCAmelCase       :       List[str]							   =self.activation(snake_case)
          return hidden_state






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






       def __init__(      self  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case = 2  ,							snake_case = 2  ,							**snake_case)						->  List[Any]:




          '''simple docstring'''
          super().__init__(**snake_case)

          _UpperCAmelCase       :       List[Any]							   =TFRegNetXLayer if config.layer_type == 'x' else TFRegNetYLayer
          _UpperCAmelCase       :       Any							   =[
              # 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       lowerCAmelCase		(      self  ,							snake_case)						->  str:




          '''simple docstring'''
          for layer_module in self.layers:
             _UpperCAmelCase       :       Optional[Any]							   =layer_module(snake_case)
          return hidden_state






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






       def __init__(      self  ,							snake_case  ,							**snake_case)						->  str:




          '''simple docstring'''
          super().__init__(**snake_case)
          _UpperCAmelCase       :       str							   =[]
          # 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'  ,							))
          _UpperCAmelCase       :       Optional[int]							   =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       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case = False  ,							snake_case = True)						->  TFBaseModelOutputWithNoAttention:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[Any]							   =() if output_hidden_states else None

          for stage_module in self.stages:
             if output_hidden_states:
                _UpperCAmelCase       :       Union[str, Any]							   =hidden_states + (hidden_state,)

             _UpperCAmelCase       :       Optional[Any]							   =stage_module(snake_case)

          if output_hidden_states:
             _UpperCAmelCase       :       Union[str, 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=snake_case  ,							hidden_states=snake_case)






@keras_serializable
class  __magic_name__							(		tf.keras.layers.Layer    ):
       UpperCAmelCase										=RegNetConfig






       def __init__(      self  ,							snake_case  ,							**snake_case)						->  int:




          '''simple docstring'''
          super().__init__(**snake_case)
          _UpperCAmelCase       :       str							   =config
          _UpperCAmelCase       :       int							   =TFRegNetEmbeddings(snake_case  ,							name='embedder')
          _UpperCAmelCase       :       int							   =TFRegNetEncoder(snake_case  ,							name='encoder')
          _UpperCAmelCase       :       Dict							   =tf.keras.layers.GlobalAveragePoolingaD(keepdims=snake_case  ,							name='pooler')






       @unpack_inputs
       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case = None  ,							snake_case = None  ,							snake_case = False  ,							)						->  TFBaseModelOutputWithPoolingAndNoAttention:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   =(
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          _UpperCAmelCase       :       Union[str, Any]							   =return_dict if return_dict is not None else self.config.use_return_dict

          _UpperCAmelCase       :       Dict							   =self.embedder(snake_case  ,							training=snake_case)

          _UpperCAmelCase       :       int							   =self.encoder(
              snake_case  ,							output_hidden_states=snake_case  ,							return_dict=snake_case  ,							training=snake_case)

          _UpperCAmelCase       :       List[str]							   =encoder_outputs[0]
          _UpperCAmelCase       :       Dict							   =self.pooler(snake_case)

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

          # Change the other hidden state outputs to NCHW as well
          if output_hidden_states:
             _UpperCAmelCase       :       Union[str, Any]							   =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  __magic_name__							(		lowerCAmelCase    ):
       UpperCAmelCase										=RegNetConfig
       UpperCAmelCase										="regnet"
       UpperCAmelCase										="pixel_values"






       @property
       def       lowerCAmelCase		(      self)						->  Any:




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






lowercase							=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							=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." ,lowerCAmelCase ,)
class  __magic_name__							(		lowerCAmelCase    ):






       def __init__(      self  ,							snake_case  ,							*snake_case  ,							**snake_case)						->  Optional[Any]:




          '''simple docstring'''
          super().__init__(snake_case  ,							*snake_case  ,							**snake_case)
          _UpperCAmelCase       :       str							   =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       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case = None  ,							snake_case = None  ,							snake_case=False  ,							)						->  Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   =(
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          _UpperCAmelCase       :       Tuple							   =return_dict if return_dict is not None else self.config.use_return_dict

          _UpperCAmelCase       :       Any							   =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    " ,lowerCAmelCase ,)
class  __magic_name__							(		lowerCAmelCase ,lowerCAmelCase    ):






       def __init__(      self  ,							snake_case  ,							*snake_case  ,							**snake_case)						->  Optional[Any]:




          '''simple docstring'''
          super().__init__(snake_case  ,							*snake_case  ,							**snake_case)
          _UpperCAmelCase       :       int							   =config.num_labels
          _UpperCAmelCase       :       Tuple							   =TFRegNetMainLayer(snake_case  ,							name='regnet')
          # classification head
          _UpperCAmelCase       :       Any							   =[
              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       lowerCAmelCase		(      self  ,							snake_case = None  ,							snake_case = None  ,							snake_case = None  ,							snake_case = None  ,							snake_case=False  ,							)						->  Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:




          '''simple docstring'''
          _UpperCAmelCase       :       Tuple							   =(
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          _UpperCAmelCase       :       Dict							   =return_dict if return_dict is not None else self.config.use_return_dict

          _UpperCAmelCase       :       Optional[int]							   =self.regnet(
              snake_case  ,							output_hidden_states=snake_case  ,							return_dict=snake_case  ,							training=snake_case)

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

          _UpperCAmelCase       :       Union[str, Any]							   =self.classifier[0](snake_case)
          _UpperCAmelCase       :       Union[str, Any]							   =self.classifier[1](snake_case)

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

          if not return_dict:
             _UpperCAmelCase       :       List[str]							   =(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'''






from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowercase							=logging.get_logger(__name__)

lowercase							={
    'facebook/levit-128S': 'https://huggingface.co/facebook/levit-128S/resolve/main/config.json',
    # See all LeViT models at https://huggingface.co/models?filter=levit
}






class  __magic_name__							(		lowerCAmelCase    ):
       UpperCAmelCase										="levit"






       def __init__(      self  ,							snake_case=2_2_4  ,							snake_case=3  ,							snake_case=3  ,							snake_case=2  ,							snake_case=1  ,							snake_case=1_6  ,							snake_case=[1_2_8, 2_5_6, 3_8_4]  ,							snake_case=[4, 8, 1_2]  ,							snake_case=[4, 4, 4]  ,							snake_case=[1_6, 1_6, 1_6]  ,							snake_case=0  ,							snake_case=[2, 2, 2]  ,							snake_case=[2, 2, 2]  ,							snake_case=0.02  ,							**snake_case  ,							)						->  Any:




          '''simple docstring'''
          super().__init__(**snake_case)
          _UpperCAmelCase       :       List[str]							   =image_size
          _UpperCAmelCase       :       str							   =num_channels
          _UpperCAmelCase       :       int							   =kernel_size
          _UpperCAmelCase       :       Any							   =stride
          _UpperCAmelCase       :       Tuple							   =padding
          _UpperCAmelCase       :       Optional[int]							   =hidden_sizes
          _UpperCAmelCase       :       List[str]							   =num_attention_heads
          _UpperCAmelCase       :       List[Any]							   =depths
          _UpperCAmelCase       :       Any							   =key_dim
          _UpperCAmelCase       :       List[Any]							   =drop_path_rate
          _UpperCAmelCase       :       int							   =patch_size
          _UpperCAmelCase       :       Tuple							   =attention_ratio
          _UpperCAmelCase       :       Any							   =mlp_ratio
          _UpperCAmelCase       :       Optional[Any]							   =initializer_range
          _UpperCAmelCase       :       str							   =[
              ['Subsample', key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
              ['Subsample', key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
          ]






class  __magic_name__							(		lowerCAmelCase    ):
       UpperCAmelCase										=version.parse("1.11"    )






       @property
       def       lowerCAmelCase		(      self)						->  Mapping[str, Mapping[int, str]]:




          '''simple docstring'''
          return OrderedDict(
              [
                  ('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'}),
              ])






       @property
       def       lowerCAmelCase		(      self)						->  float:




          '''simple docstring'''
          return 1E-4

