


'''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:						str  		=   logging.get_logger(__name__)

# General docstring
__lowercase:						str  		=   """ResNetConfig"""

# Base docstring
__lowercase:						str  		=   """microsoft/resnet-50"""
__lowercase:						Union[str, Any]  		=   [1, 2_048, 7, 7]

# Image classification docstring
__lowercase:						List[Any]  		=   """microsoft/resnet-50"""
__lowercase:						Any  		=   """tiger cat"""

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







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

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





    """simple docstring"""



    super().__init__(**__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  kernel_size // 2
    UpperCamelCase__         =  tf.keras.layers.ConvaD(
        __SCREAMING_SNAKE_CASE,    kernel_size=__SCREAMING_SNAKE_CASE,    strides=__SCREAMING_SNAKE_CASE,    padding="valid",    use_bias=__SCREAMING_SNAKE_CASE,    name="convolution"    )
    # Use same default momentum and epsilon as PyTorch equivalent
    UpperCamelCase__         =  tf.keras.layers.BatchNormalization(epsilon=1e-5,    momentum=0.9,    name="normalization"    )
    UpperCamelCase__         =  ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("linear"    )

 def 						lowercase_  (       self       :		Dict,    a_       :		int    ):





    """simple docstring"""



    UpperCamelCase__         =  UpperCamelCase__         =  (self.pad_value, self.pad_value)
    UpperCamelCase__         =  tf.pad(__SCREAMING_SNAKE_CASE,    [(0, 0), height_pad, width_pad, (0, 0)]    )
    UpperCamelCase__         =  self.conv(__SCREAMING_SNAKE_CASE    )
    return hidden_state

 def 						lowercase_  (       self       :		Optional[int],    a_       :		List[Any],    a_       :		List[str] = False    ):





    """simple docstring"""



    UpperCamelCase__         =  self.convolution(__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  self.normalization(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  self.activation(__SCREAMING_SNAKE_CASE    )
    return hidden_state







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

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





    """simple docstring"""



    super().__init__(**__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  TFResNetConvLayer(
        config.embedding_size,    kernel_size=7,    stride=2,    activation=config.hidden_act,    name="embedder",    )
    UpperCamelCase__         =  tf.keras.layers.MaxPoolaD(pool_size=3,    strides=2,    padding="valid",    name="pooler"    )
    UpperCamelCase__         =  config.num_channels

 def 						lowercase_  (       self       :		str,    a_       :		List[Any],    a_       :		List[Any] = False    ):





    """simple docstring"""



    UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__         =  shape_list(__SCREAMING_SNAKE_CASE    )
    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."    )
    UpperCamelCase__         =  pixel_values
    UpperCamelCase__         =  self.embedder(__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  tf.pad(__SCREAMING_SNAKE_CASE,    [[0, 0], [1, 1], [1, 1], [0, 0]]    )
    UpperCamelCase__         =  self.pooler(__SCREAMING_SNAKE_CASE    )
    return hidden_state







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

 def __init__(       self       :		List[str],    a_       :		int,    a_       :		Any = 2,    **a_       :		str    ):





    """simple docstring"""



    super().__init__(**__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  tf.keras.layers.ConvaD(
        __SCREAMING_SNAKE_CASE,    kernel_size=1,    strides=__SCREAMING_SNAKE_CASE,    use_bias=__SCREAMING_SNAKE_CASE,    name="convolution"    )
    # Use same default momentum and epsilon as PyTorch equivalent
    UpperCamelCase__         =  tf.keras.layers.BatchNormalization(epsilon=1e-5,    momentum=0.9,    name="normalization"    )

 def 						lowercase_  (       self       :		Union[str, Any],    a_       :		Tuple,    a_       :		str = False    ):





    """simple docstring"""



    UpperCamelCase__         =  x
    UpperCamelCase__         =  self.convolution(__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  self.normalization(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )
    return hidden_state







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

 def __init__(       self       :		List[str],    a_       :		Union[str, Any],    a_       :		List[str],    a_       :		Any = 1,    a_       :		List[Any] = "relu",    **a_       :		Dict    ):





    """simple docstring"""



    super().__init__(**__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  in_channels != out_channels or stride != 1
    UpperCamelCase__         =  TFResNetConvLayer(__SCREAMING_SNAKE_CASE,    stride=__SCREAMING_SNAKE_CASE,    name="layer.0"    )
    UpperCamelCase__         =  TFResNetConvLayer(__SCREAMING_SNAKE_CASE,    activation=__SCREAMING_SNAKE_CASE,    name="layer.1"    )
    UpperCamelCase__         =  (
        TFResNetShortCut(__SCREAMING_SNAKE_CASE,    stride=__SCREAMING_SNAKE_CASE,    name="shortcut"    )
        if should_apply_shortcut
        else tf.keras.layers.Activation("linear",    name="shortcut"    )
    )
    UpperCamelCase__         =  ACTaFN[activation]

 def 						lowercase_  (       self       :		Tuple,    a_       :		List[str],    a_       :		Optional[int] = False    ):





    """simple docstring"""



    UpperCamelCase__         =  hidden_state
    UpperCamelCase__         =  self.conva(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  self.conva(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  self.shortcut(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )
    hidden_state += residual
    UpperCamelCase__         =  self.activation(__SCREAMING_SNAKE_CASE    )
    return hidden_state







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

 def __init__(       self       :		Tuple,    a_       :		List[Any],    a_       :		List[str],    a_       :		Optional[int] = 1,    a_       :		Optional[Any] = "relu",    a_       :		List[str] = 4,    **a_       :		List[str],    ):





    """simple docstring"""



    super().__init__(**__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  in_channels != out_channels or stride != 1
    UpperCamelCase__         =  out_channels // reduction
    UpperCamelCase__         =  TFResNetConvLayer(__SCREAMING_SNAKE_CASE,    kernel_size=1,    name="layer.0"    )
    UpperCamelCase__         =  TFResNetConvLayer(__SCREAMING_SNAKE_CASE,    stride=__SCREAMING_SNAKE_CASE,    name="layer.1"    )
    UpperCamelCase__         =  TFResNetConvLayer(__SCREAMING_SNAKE_CASE,    kernel_size=1,    activation=__SCREAMING_SNAKE_CASE,    name="layer.2"    )
    UpperCamelCase__         =  (
        TFResNetShortCut(__SCREAMING_SNAKE_CASE,    stride=__SCREAMING_SNAKE_CASE,    name="shortcut"    )
        if should_apply_shortcut
        else tf.keras.layers.Activation("linear",    name="shortcut"    )
    )
    UpperCamelCase__         =  ACTaFN[activation]

 def 						lowercase_  (       self       :		List[Any],    a_       :		str,    a_       :		Union[str, Any] = False    ):





    """simple docstring"""



    UpperCamelCase__         =  hidden_state
    UpperCamelCase__         =  self.conva(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  self.conva(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  self.conva(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  self.shortcut(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )
    hidden_state += residual
    UpperCamelCase__         =  self.activation(__SCREAMING_SNAKE_CASE    )
    return hidden_state







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

 def __init__(       self       :		Any,    a_       :		Dict,    a_       :		int,    a_       :		str,    a_       :		Optional[int] = 2,    a_       :		Any = 2,    **a_       :		Any    ):





    """simple docstring"""



    super().__init__(**__SCREAMING_SNAKE_CASE    )

    UpperCamelCase__         =  TFResNetBottleNeckLayer if config.layer_type == "bottleneck" else TFResNetBasicLayer

    UpperCamelCase__         =  [layer(__SCREAMING_SNAKE_CASE,    __SCREAMING_SNAKE_CASE,    stride=__SCREAMING_SNAKE_CASE,    activation=config.hidden_act,    name="layers.0"    )]
    layers += [
        layer(__SCREAMING_SNAKE_CASE,    __SCREAMING_SNAKE_CASE,    activation=config.hidden_act,    name=f'layers.{i + 1}'    )
        for i in range(depth - 1    )
    ]
    UpperCamelCase__         =  layers

 def 						lowercase_  (       self       :		Tuple,    a_       :		int,    a_       :		int = False    ):





    """simple docstring"""



    for layer in self.stage_layers:
       UpperCamelCase__         =  layer(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )
    return hidden_state







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

 def __init__(       self       :		Dict,    a_       :		Optional[Any],    **a_       :		str    ):





    """simple docstring"""



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

 def 						lowercase_  (       self       :		Tuple,    a_       :		Any,    a_       :		Tuple = False,    a_       :		Optional[int] = True,    a_       :		Tuple = False,    ):





    """simple docstring"""



    UpperCamelCase__         =  () if output_hidden_states else None

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

       UpperCamelCase__         =  stage_module(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )

    if output_hidden_states:
       UpperCamelCase__         =  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=__SCREAMING_SNAKE_CASE,    hidden_states=__SCREAMING_SNAKE_CASE    )







class        UpperCAmelCase       (				__lowerCAmelCase):
 _lowerCamelCase       :    Dict			    =    ResNetConfig
 _lowerCamelCase       :    Any			    =    """resnet"""
 _lowerCamelCase       :    str			    =    """pixel_values"""

 @property
 def 						lowercase_  (       self       :		Dict    ):





    """simple docstring"""



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







__lowercase:						Any  		=   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:						Tuple  		=   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        UpperCAmelCase       (				tf.keras.layers.Layer):
 _lowerCamelCase       :    Tuple			    =    ResNetConfig

 def __init__(       self       :		Any,    a_       :		str,    **a_       :		Tuple    ):





    """simple docstring"""



    super().__init__(**__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  config
    UpperCamelCase__         =  TFResNetEmbeddings(__SCREAMING_SNAKE_CASE,    name="embedder"    )
    UpperCamelCase__         =  TFResNetEncoder(__SCREAMING_SNAKE_CASE,    name="encoder"    )
    UpperCamelCase__         =  tf.keras.layers.GlobalAveragePoolingaD(keepdims=__SCREAMING_SNAKE_CASE    )

 @unpack_inputs
 def 						lowercase_  (       self       :		List[str],    a_       :		int,    a_       :		List[str] = None,    a_       :		Optional[int] = None,    a_       :		List[Any] = False,    ):





    """simple docstring"""



    UpperCamelCase__         =  (
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    UpperCamelCase__         =  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)
    UpperCamelCase__         =  tf.transpose(__SCREAMING_SNAKE_CASE,    perm=[0, 2, 3, 1]    )
    UpperCamelCase__         =  self.embedder(__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )

    UpperCamelCase__         =  self.encoder(
        __SCREAMING_SNAKE_CASE,    output_hidden_states=__SCREAMING_SNAKE_CASE,    return_dict=__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )

    UpperCamelCase__         =  encoder_outputs[0]

    UpperCamelCase__         =  self.pooler(__SCREAMING_SNAKE_CASE    )

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

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

    UpperCamelCase__         =  hidden_states if output_hidden_states else None

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







@add_start_docstrings(
    'The bare ResNet model outputting raw features without any specific head on top.'						,	__lowerCAmelCase						,	)
class        UpperCAmelCase       (				__lowerCAmelCase):

 def __init__(       self       :		Union[str, Any],    a_       :		Tuple,    **a_       :		Dict    ):





    """simple docstring"""



    super().__init__(__SCREAMING_SNAKE_CASE,    **__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  TFResNetMainLayer(config=__SCREAMING_SNAKE_CASE,    name="resnet"    )

 @add_start_docstrings_to_model_forward(__SCREAMING_SNAKE_CASE    )
 @add_code_sample_docstrings(
     checkpoint=_CHECKPOINT_FOR_DOC,    output_type=__SCREAMING_SNAKE_CASE,    config_class=_CONFIG_FOR_DOC,    modality="vision",    expected_output=_EXPECTED_OUTPUT_SHAPE,    )
 @unpack_inputs
 def 						lowercase_  (       self       :		Tuple,    a_       :		int,    a_       :		Dict = None,    a_       :		int = None,    a_       :		List[Any] = False,    ):





    """simple docstring"""



    UpperCamelCase__         =  (
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    UpperCamelCase__         =  return_dict if return_dict is not None else self.config.use_return_dict

    UpperCamelCase__         =  self.resnet(
        pixel_values=__SCREAMING_SNAKE_CASE,    output_hidden_states=__SCREAMING_SNAKE_CASE,    return_dict=__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE,    )
    return resnet_outputs







@add_start_docstrings(
    '\n    ResNet 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        UpperCAmelCase       (				__lowerCAmelCase						,	__lowerCAmelCase):

 def __init__(       self       :		Tuple,    a_       :		Optional[Any],    **a_       :		str    ):





    """simple docstring"""



    super().__init__(__SCREAMING_SNAKE_CASE,    **__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  config.num_labels
    UpperCamelCase__         =  TFResNetMainLayer(__SCREAMING_SNAKE_CASE,    name="resnet"    )
    # classification head
    UpperCamelCase__         =  (
        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 						lowercase_  (       self       :		str,    a_       :		Optional[int]    ):





    """simple docstring"""



    UpperCamelCase__         =  tf.keras.layers.Flatten()(__SCREAMING_SNAKE_CASE    )
    UpperCamelCase__         =  self.classifier_layer(__SCREAMING_SNAKE_CASE    )
    return logits

 @add_start_docstrings_to_model_forward(__SCREAMING_SNAKE_CASE    )
 @add_code_sample_docstrings(
     checkpoint=_IMAGE_CLASS_CHECKPOINT,    output_type=__SCREAMING_SNAKE_CASE,    config_class=_CONFIG_FOR_DOC,    expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,    )
 @unpack_inputs
 def 						lowercase_  (       self       :		List[Any],    a_       :		str = None,    a_       :		Dict = None,    a_       :		str = None,    a_       :		Dict = None,    a_       :		List[str] = False,    ):





    """simple docstring"""



    UpperCamelCase__         =  return_dict if return_dict is not None else self.config.use_return_dict

    UpperCamelCase__         =  self.resnet(
        __SCREAMING_SNAKE_CASE,    output_hidden_states=__SCREAMING_SNAKE_CASE,    return_dict=__SCREAMING_SNAKE_CASE,    training=__SCREAMING_SNAKE_CASE    )

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

    UpperCamelCase__         =  self.classifier(__SCREAMING_SNAKE_CASE    )

    UpperCamelCase__         =  None if labels is None else self.hf_compute_loss(__SCREAMING_SNAKE_CASE,    __SCREAMING_SNAKE_CASE    )

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

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






"""simple docstring"""


import warnings

from ...utils import logging
from .image_processing_vilt import ViltImageProcessor


a__    :   str    						= logging.get_logger(__name__)



class    lowercase_ (   __lowerCAmelCase							):

       def __init__(	self      ,    *__SCREAMING_SNAKE_CASE      ,    **__SCREAMING_SNAKE_CASE   ):






        '''simple docstring'''

        warnings.warn(
            "The class ViltFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
            " use ViltImageProcessor instead."      ,    __SCREAMING_SNAKE_CASE      ,    )
        super().__init__(*__SCREAMING_SNAKE_CASE      ,    **__SCREAMING_SNAKE_CASE   )

