


"""simple docstring"""







from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFImageClassifierOutputWithNoAttention,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_resnet import ResNetConfig


__lowercase          = logging.get_logger(__name__)

# General docstring
__lowercase          = """ResNetConfig"""

# Base docstring
__lowercase          = """microsoft/resnet-50"""
__lowercase          = [1, 2048, 7, 7]

# Image classification docstring
__lowercase          = """microsoft/resnet-50"""
__lowercase          = """tiger cat"""

__lowercase          = [
    """microsoft/resnet-50""",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]






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







    """simple docstring"""




    def __init__(   self							:       Dict							,					__UpperCAmelCase							:       int							,					__UpperCAmelCase							:       int = 3							,					__UpperCAmelCase							:       int = 1							,					__UpperCAmelCase							:       str = "relu"							,					**__UpperCAmelCase							:       int):
      super().__init__(**__UpperCAmelCase)
      a		:						List[str] 				=     kernel_size // 2
      a		:						Tuple 				=     tf.keras.layers.ConvaD(
          __UpperCAmelCase							,					kernel_size=__UpperCAmelCase							,					strides=__UpperCAmelCase							,					padding="valid"							,					use_bias=__UpperCAmelCase							,					name="convolution")
      # Use same default momentum and epsilon as PyTorch equivalent
      a		:						Optional[int] 				=     tf.keras.layers.BatchNormalization(epsilon=1e-5							,					momentum=0.9							,					name="normalization")
      a		:						Optional[Any] 				=     ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("linear")






    def       __snake_case				(   self							:       int							,					__UpperCAmelCase							:       tf.Tensor):
      # Pad to match that done in the PyTorch Conv2D model
      a		:						Tuple 				=     (self.pad_value, self.pad_value)
      a		:						Dict 				=     tf.pad(__UpperCAmelCase							,					[(0, 0), height_pad, width_pad, (0, 0)])
      a		:						Tuple 				=     self.conv(__UpperCAmelCase)
      return hidden_state






    def       __snake_case				(   self							:       Optional[int]							,					__UpperCAmelCase							:       tf.Tensor							,					__UpperCAmelCase							:       bool = False):
      a		:						Dict 				=     self.convolution(__UpperCAmelCase)
      a		:						int 				=     self.normalization(__UpperCAmelCase							,					training=__UpperCAmelCase)
      a		:						Any 				=     self.activation(__UpperCAmelCase)
      return hidden_state






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







    """simple docstring"""




    def __init__(   self							:       str							,					__UpperCAmelCase							:       ResNetConfig							,					**__UpperCAmelCase							:       List[Any]):
      super().__init__(**__UpperCAmelCase)
      a		:						Dict 				=     TFResNetConvLayer(
          config.embedding_size							,					kernel_size=7							,					stride=2							,					activation=config.hidden_act							,					name="embedder"							,					)
      a		:						str 				=     tf.keras.layers.MaxPoolaD(pool_size=3							,					strides=2							,					padding="valid"							,					name="pooler")
      a		:						Union[str, Any] 				=     config.num_channels






    def       __snake_case				(   self							:       str							,					__UpperCAmelCase							:       tf.Tensor							,					__UpperCAmelCase							:       bool = False):
      a				,						a				,						a				,						a		:						Optional[int] 				=     shape_list(__UpperCAmelCase)
      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.")
      a		:						List[str] 				=     pixel_values
      a		:						int 				=     self.embedder(__UpperCAmelCase)
      a		:						Union[str, Any] 				=     tf.pad(__UpperCAmelCase							,					[[0, 0], [1, 1], [1, 1], [0, 0]])
      a		:						Optional[int] 				=     self.pooler(__UpperCAmelCase)
      return hidden_state






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







    """simple docstring"""




    def __init__(   self							:       str							,					__UpperCAmelCase							:       int							,					__UpperCAmelCase							:       int = 2							,					**__UpperCAmelCase							:       Optional[Any]):
      super().__init__(**__UpperCAmelCase)
      a		:						Any 				=     tf.keras.layers.ConvaD(
          __UpperCAmelCase							,					kernel_size=1							,					strides=__UpperCAmelCase							,					use_bias=__UpperCAmelCase							,					name="convolution")
      # Use same default momentum and epsilon as PyTorch equivalent
      a		:						Optional[Any] 				=     tf.keras.layers.BatchNormalization(epsilon=1e-5							,					momentum=0.9							,					name="normalization")






    def       __snake_case				(   self							:       Union[str, Any]							,					__UpperCAmelCase							:       tf.Tensor							,					__UpperCAmelCase							:       bool = False):
      a		:						int 				=     x
      a		:						Dict 				=     self.convolution(__UpperCAmelCase)
      a		:						Optional[Any] 				=     self.normalization(__UpperCAmelCase							,					training=__UpperCAmelCase)
      return hidden_state






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







    """simple docstring"""




    def __init__(   self							:       Union[str, Any]							,					__UpperCAmelCase							:       int							,					__UpperCAmelCase							:       int							,					__UpperCAmelCase							:       int = 1							,					__UpperCAmelCase							:       str = "relu"							,					**__UpperCAmelCase							:       Tuple):
      super().__init__(**__UpperCAmelCase)
      a		:						List[str] 				=     in_channels != out_channels or stride != 1
      a		:						int 				=     TFResNetConvLayer(__UpperCAmelCase							,					stride=__UpperCAmelCase							,					name="layer.0")
      a		:						Any 				=     TFResNetConvLayer(__UpperCAmelCase							,					activation=__UpperCAmelCase							,					name="layer.1")
      a		:						Any 				=     (
          TFResNetShortCut(__UpperCAmelCase							,					stride=__UpperCAmelCase							,					name="shortcut")
          if should_apply_shortcut
          else tf.keras.layers.Activation("linear"							,					name="shortcut")
      )
      a		:						int 				=     ACTaFN[activation]






    def       __snake_case				(   self							:       str							,					__UpperCAmelCase							:       tf.Tensor							,					__UpperCAmelCase							:       bool = False):
      a		:						Any 				=     hidden_state
      a		:						Tuple 				=     self.conva(__UpperCAmelCase							,					training=__UpperCAmelCase)
      a		:						str 				=     self.conva(__UpperCAmelCase							,					training=__UpperCAmelCase)
      a		:						Optional[Any] 				=     self.shortcut(__UpperCAmelCase							,					training=__UpperCAmelCase)
      hidden_state += residual
      a		:						Optional[Any] 				=     self.activation(__UpperCAmelCase)
      return hidden_state






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







    """simple docstring"""




    def __init__(   self							:       Optional[Any]							,					__UpperCAmelCase							:       int							,					__UpperCAmelCase							:       int							,					__UpperCAmelCase							:       int = 1							,					__UpperCAmelCase							:       str = "relu"							,					__UpperCAmelCase							:       int = 4							,					**__UpperCAmelCase							:       List[Any]							,					):
      super().__init__(**__UpperCAmelCase)
      a		:						List[Any] 				=     in_channels != out_channels or stride != 1
      a		:						Tuple 				=     out_channels // reduction
      a		:						Dict 				=     TFResNetConvLayer(__UpperCAmelCase							,					kernel_size=1							,					name="layer.0")
      a		:						str 				=     TFResNetConvLayer(__UpperCAmelCase							,					stride=__UpperCAmelCase							,					name="layer.1")
      a		:						Dict 				=     TFResNetConvLayer(__UpperCAmelCase							,					kernel_size=1							,					activation=__UpperCAmelCase							,					name="layer.2")
      a		:						Tuple 				=     (
          TFResNetShortCut(__UpperCAmelCase							,					stride=__UpperCAmelCase							,					name="shortcut")
          if should_apply_shortcut
          else tf.keras.layers.Activation("linear"							,					name="shortcut")
      )
      a		:						Dict 				=     ACTaFN[activation]






    def       __snake_case				(   self							:       Optional[int]							,					__UpperCAmelCase							:       tf.Tensor							,					__UpperCAmelCase							:       bool = False):
      a		:						List[Any] 				=     hidden_state
      a		:						str 				=     self.conva(__UpperCAmelCase							,					training=__UpperCAmelCase)
      a		:						Any 				=     self.conva(__UpperCAmelCase							,					training=__UpperCAmelCase)
      a		:						int 				=     self.conva(__UpperCAmelCase							,					training=__UpperCAmelCase)
      a		:						int 				=     self.shortcut(__UpperCAmelCase							,					training=__UpperCAmelCase)
      hidden_state += residual
      a		:						int 				=     self.activation(__UpperCAmelCase)
      return hidden_state






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







    """simple docstring"""




    def __init__(   self							:       Tuple							,					__UpperCAmelCase							:       ResNetConfig							,					__UpperCAmelCase							:       int							,					__UpperCAmelCase							:       int							,					__UpperCAmelCase							:       int = 2							,					__UpperCAmelCase							:       int = 2							,					**__UpperCAmelCase							:       Union[str, Any]):
      super().__init__(**__UpperCAmelCase)

      a		:						Any 				=     TFResNetBottleNeckLayer if config.layer_type == "bottleneck" else TFResNetBasicLayer

      a		:						Any 				=     [layer(__UpperCAmelCase							,					__UpperCAmelCase							,					stride=__UpperCAmelCase							,					activation=config.hidden_act							,					name="layers.0")]
      layers += [
          layer(__UpperCAmelCase							,					__UpperCAmelCase							,					activation=config.hidden_act							,					name=f'''layers.{i + 1}''')
          for i in range(depth - 1)
      ]
      a		:						Union[str, Any] 				=     layers






    def       __snake_case				(   self							:       str							,					__UpperCAmelCase							:       tf.Tensor							,					__UpperCAmelCase							:       bool = False):
      for layer in self.stage_layers:
        a		:						Optional[int] 				=     layer(__UpperCAmelCase							,					training=__UpperCAmelCase)
      return hidden_state






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







    """simple docstring"""




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






    def       __snake_case				(   self							:       Dict							,					__UpperCAmelCase							:       tf.Tensor							,					__UpperCAmelCase							:       bool = False							,					__UpperCAmelCase							:       bool = True							,					__UpperCAmelCase							:       bool = False							,					):
      a		:						int 				=     () if output_hidden_states else None

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

        a		:						str 				=     stage_module(__UpperCAmelCase							,					training=__UpperCAmelCase)

      if output_hidden_states:
        a		:						Dict 				=     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)






class 		_A    (    _a    ):







    """simple docstring"""




    UpperCAmelCase			:				List[Any]      =			ResNetConfig
    UpperCAmelCase			:				Union[str, Any]      =			"""resnet"""
    UpperCAmelCase			:				List[Any]      =			"""pixel_values"""






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






__lowercase          = R"""
    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.

    Parameters:
        config ([`ResNetConfig`]): 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.
"""


__lowercase          = R"""
    Args:
        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__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.
"""






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







    """simple docstring"""




    UpperCAmelCase			:				Any      =			ResNetConfig






    def __init__(   self							:       Optional[int]							,					__UpperCAmelCase							:       ResNetConfig							,					**__UpperCAmelCase							:       List[Any]):
      super().__init__(**__UpperCAmelCase)
      a		:						Dict 				=     config
      a		:						Union[str, Any] 				=     TFResNetEmbeddings(__UpperCAmelCase							,					name="embedder")
      a		:						int 				=     TFResNetEncoder(__UpperCAmelCase							,					name="encoder")
      a		:						Union[str, Any] 				=     tf.keras.layers.GlobalAveragePoolingaD(keepdims=__UpperCAmelCase)






    @unpack_inputs
    def       __snake_case				(   self							:       Union[str, Any]							,					__UpperCAmelCase							:       tf.Tensor							,					__UpperCAmelCase							:       Optional[bool] = None							,					__UpperCAmelCase							:       Optional[bool] = None							,					__UpperCAmelCase							:       bool = False							,					):
      a		:						int 				=     (
          output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
      )
      a		:						Optional[Any] 				=     return_dict if return_dict is not None else self.config.use_return_dict

      # TF 2.0 image layers can't use NCHW format when running on CPU.
      # We transpose to NHWC format and then transpose back after the full forward pass.
      # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
      a		:						str 				=     tf.transpose(__UpperCAmelCase							,					perm=[0, 2, 3, 1])
      a		:						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		:						List[str] 				=     self.pooler(__UpperCAmelCase)

      # Transpose all the outputs to the NCHW format
      # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
      a		:						Any 				=     tf.transpose(__UpperCAmelCase							,					(0, 3, 1, 2))
      a		:						List[Any] 				=     tf.transpose(__UpperCAmelCase							,					(0, 3, 1, 2))
      a		:						List[Any] 				=     ()
      for hidden_state in encoder_outputs[1:]:
        a		:						List[Any] 				=     hidden_states + tuple(tf.transpose(__UpperCAmelCase							,					(0, 3, 1, 2)) for h in hidden_state)

      if not return_dict:
        return (last_hidden_state, pooled_output) + hidden_states

      a		:						Any 				=     hidden_states if output_hidden_states else None

      return TFBaseModelOutputWithPoolingAndNoAttention(
          last_hidden_state=__UpperCAmelCase							,					pooler_output=__UpperCAmelCase							,					hidden_states=__UpperCAmelCase							,					)






@add_start_docstrings(
    """The bare ResNet model outputting raw features without any specific head on top."""       ,_a       ,)
class 		_A    (    _a    ):







    """simple docstring"""




    def __init__(   self							:       Dict							,					__UpperCAmelCase							:       ResNetConfig							,					**__UpperCAmelCase							:       Dict):
      super().__init__(__UpperCAmelCase							,					**__UpperCAmelCase)
      a		:						Tuple 				=     TFResNetMainLayer(config=__UpperCAmelCase							,					name="resnet")






    @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							,					)
    @unpack_inputs
    def       __snake_case				(   self							:       str							,					__UpperCAmelCase							:       tf.Tensor							,					__UpperCAmelCase							:       Optional[bool] = None							,					__UpperCAmelCase							:       Optional[bool] = None							,					__UpperCAmelCase							:       bool = False							,					):
      a		:						List[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		:						List[str] 				=     self.resnet(
          pixel_values=__UpperCAmelCase							,					output_hidden_states=__UpperCAmelCase							,					return_dict=__UpperCAmelCase							,					training=__UpperCAmelCase							,					)
      return resnet_outputs






@add_start_docstrings(
    """
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    """       ,_a       ,)
class 		_A    (    _a       ,_a    ):







    """simple docstring"""




    def __init__(   self							:       List[Any]							,					__UpperCAmelCase							:       ResNetConfig							,					**__UpperCAmelCase							:       List[Any]):
      super().__init__(__UpperCAmelCase							,					**__UpperCAmelCase)
      a		:						Union[str, Any] 				=     config.num_labels
      a		:						Union[str, Any] 				=     TFResNetMainLayer(__UpperCAmelCase							,					name="resnet")
      # classification head
      a		:						List[str] 				=     (
          tf.keras.layers.Dense(config.num_labels							,					name="classifier.1")
          if config.num_labels > 0
          else tf.keras.layers.Activation("linear"							,					name="classifier.1")
      )






    def       __snake_case				(   self							:       Union[str, Any]							,					__UpperCAmelCase							:       tf.Tensor):
      a		:						List[str] 				=     tf.keras.layers.Flatten()(__UpperCAmelCase)
      a		:						Tuple 				=     self.classifier_layer(__UpperCAmelCase)
      return logits






    @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							,					)
    @unpack_inputs
    def       __snake_case				(   self							:       Tuple							,					__UpperCAmelCase							:       tf.Tensor = None							,					__UpperCAmelCase							:       tf.Tensor = None							,					__UpperCAmelCase							:       bool = None							,					__UpperCAmelCase							:       bool = None							,					__UpperCAmelCase							:       bool = False							,					):
      a		:						Tuple 				=     return_dict if return_dict is not None else self.config.use_return_dict

      a		:						int 				=     self.resnet(
          __UpperCAmelCase							,					output_hidden_states=__UpperCAmelCase							,					return_dict=__UpperCAmelCase							,					training=__UpperCAmelCase)

      a		:						str 				=     outputs.pooler_output if return_dict else outputs[1]

      a		:						Tuple 				=     self.classifier(__UpperCAmelCase)

      a		:						Optional[Any] 				=     None if labels is None else self.hf_compute_loss(__UpperCAmelCase							,					__UpperCAmelCase)

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

      return TFImageClassifierOutputWithNoAttention(loss=__UpperCAmelCase							,					logits=__UpperCAmelCase							,					hidden_states=outputs.hidden_states)










"""simple docstring"""







import warnings

from ...utils import logging
from .image_processing_levit import LevitImageProcessor


__lowercase          = logging.get_logger(__name__)






class 		_A    (    _a    ):







    """simple docstring"""




    def __init__(   self							:       Optional[int]							,					*__UpperCAmelCase							:       Dict							,					**__UpperCAmelCase							:       Any):
      warnings.warn(
          "The class LevitFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
          " use LevitImageProcessor instead."							,					__UpperCAmelCase							,					)
      super().__init__(*__UpperCAmelCase							,					**__UpperCAmelCase)






