






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


UpperCamelCase_						:					Optional[int]      				=			logging.get_logger(__name__)

# General docstring
UpperCamelCase_						:					Dict      				=			"""ResNetConfig"""

# Base docstring
UpperCamelCase_						:					int      				=			"""microsoft/resnet-50"""
UpperCamelCase_						:					Optional[int]      				=			[1, 2048, 7, 7]

# Image classification docstring
UpperCamelCase_						:					int      				=			"""microsoft/resnet-50"""
UpperCamelCase_						:					Dict      				=			"""tiger cat"""

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


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


      def __init__(     self					:						Tuple				, a					:						int				, a					:						int = 3				, a					:						int = 1				, a					:						str = "relu"				, **a					:						Tuple    ):





            """simple docstring"""







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


      def _UpperCamelCase    (     self					:						Dict				, a					:						tf.Tensor    ):





            """simple docstring"""







            __snake_case   :     List[Any]	      =(self.pad_value, self.pad_value)
            __snake_case   :     str	      =tf.pad(a				, [(0, 0), height_pad, width_pad, (0, 0)]    )
            __snake_case   :     str	      =self.conv(a    )
            return hidden_state







      def _UpperCamelCase    (     self					:						Union[str, Any]				, a					:						tf.Tensor				, a					:						bool = False    ):





            """simple docstring"""







            __snake_case   :     Optional[Any]	      =self.convolution(a    )
            __snake_case   :     int	      =self.normalization(a				, training=a    )
            __snake_case   :     List[Any]	      =self.activation(a    )
            return hidden_state



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


      def __init__(     self					:						int				, a					:						ResNetConfig				, **a					:						List[Any]    ):





            """simple docstring"""







            super().__init__(**a    )
            __snake_case   :     Optional[Any]	      =TFResNetConvLayer(
                config.embedding_size				, kernel_size=7				, stride=2				, activation=config.hidden_act				, name='''embedder'''				, )
            __snake_case   :     List[Any]	      =tf.keras.layers.MaxPoolaD(pool_size=3				, strides=2				, padding='''valid'''				, name='''pooler'''    )
            __snake_case   :     Tuple	      =config.num_channels







      def _UpperCamelCase    (     self					:						Optional[Any]				, a					:						tf.Tensor				, a					:						bool = False    ):





            """simple docstring"""







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



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


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





            """simple docstring"""







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







      def _UpperCamelCase    (     self					:						List[Any]				, a					:						tf.Tensor				, a					:						bool = False    ):





            """simple docstring"""







            __snake_case   :     Optional[int]	      =x
            __snake_case   :     Dict	      =self.convolution(a    )
            __snake_case   :     Union[str, Any]	      =self.normalization(a				, training=a    )
            return hidden_state



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


      def __init__(     self					:						Any				, a					:						int				, a					:						int				, a					:						int = 1				, a					:						str = "relu"				, **a					:						List[Any]    ):





            """simple docstring"""







            super().__init__(**a    )
            __snake_case   :     Any	      =in_channels != out_channels or stride != 1
            __snake_case   :     List[Any]	      =TFResNetConvLayer(a				, stride=a				, name='''layer.0'''    )
            __snake_case   :     List[Any]	      =TFResNetConvLayer(a				, activation=a				, name='''layer.1'''    )
            __snake_case   :     List[Any]	      =(
                TFResNetShortCut(a				, stride=a				, name='''shortcut'''    )
                if should_apply_shortcut
                else tf.keras.layers.Activation('''linear'''				, name='''shortcut'''    )
            )
            __snake_case   :     Any	      =ACTaFN[activation]







      def _UpperCamelCase    (     self					:						Union[str, Any]				, a					:						tf.Tensor				, a					:						bool = False    ):





            """simple docstring"""







            __snake_case   :     Optional[int]	      =hidden_state
            __snake_case   :     Dict	      =self.conva(a				, training=a    )
            __snake_case   :     Optional[int]	      =self.conva(a				, training=a    )
            __snake_case   :     Dict	      =self.shortcut(a				, training=a    )
            hidden_state += residual
            __snake_case   :     Tuple	      =self.activation(a    )
            return hidden_state



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


      def __init__(     self					:						Optional[Any]				, a					:						int				, a					:						int				, a					:						int = 1				, a					:						str = "relu"				, a					:						int = 4				, **a					:						Dict				, ):





            """simple docstring"""







            super().__init__(**a    )
            __snake_case   :     Tuple	      =in_channels != out_channels or stride != 1
            __snake_case   :     Union[str, Any]	      =out_channels // reduction
            __snake_case   :     int	      =TFResNetConvLayer(a				, kernel_size=1				, name='''layer.0'''    )
            __snake_case   :     List[str]	      =TFResNetConvLayer(a				, stride=a				, name='''layer.1'''    )
            __snake_case   :     Union[str, Any]	      =TFResNetConvLayer(a				, kernel_size=1				, activation=a				, name='''layer.2'''    )
            __snake_case   :     str	      =(
                TFResNetShortCut(a				, stride=a				, name='''shortcut'''    )
                if should_apply_shortcut
                else tf.keras.layers.Activation('''linear'''				, name='''shortcut'''    )
            )
            __snake_case   :     List[Any]	      =ACTaFN[activation]







      def _UpperCamelCase    (     self					:						Union[str, Any]				, a					:						tf.Tensor				, a					:						bool = False    ):





            """simple docstring"""







            __snake_case   :     Union[str, Any]	      =hidden_state
            __snake_case   :     Any	      =self.conva(a				, training=a    )
            __snake_case   :     Union[str, Any]	      =self.conva(a				, training=a    )
            __snake_case   :     str	      =self.conva(a				, training=a    )
            __snake_case   :     int	      =self.shortcut(a				, training=a    )
            hidden_state += residual
            __snake_case   :     Optional[int]	      =self.activation(a    )
            return hidden_state



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


      def __init__(     self					:						List[Any]				, a					:						ResNetConfig				, a					:						int				, a					:						int				, a					:						int = 2				, a					:						int = 2				, **a					:						Any    ):





            """simple docstring"""







            super().__init__(**a    )

            __snake_case   :     Union[str, Any]	      =TFResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else TFResNetBasicLayer

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







      def _UpperCamelCase    (     self					:						Union[str, Any]				, a					:						tf.Tensor				, a					:						bool = False    ):





            """simple docstring"""







            for layer in self.stage_layers:
                  __snake_case   :     int	      =layer(a				, training=a    )
            return hidden_state



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


      def __init__(     self					:						str				, a					:						ResNetConfig				, **a					:						Any    ):





            """simple docstring"""







            super().__init__(**a    )
            # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
            __snake_case   :     Dict	      =[
                TFResNetStage(
                    a				, 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(a				, a				, a				, depth=a				, name=f'''stages.{i + 1}'''    )    )







      def _UpperCamelCase    (     self					:						Union[str, Any]				, a					:						tf.Tensor				, a					:						bool = False				, a					:						bool = True				, a					:						bool = False				, ):





            """simple docstring"""







            __snake_case   :     Tuple	      =() if output_hidden_states else None

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

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

            if output_hidden_states:
                  __snake_case   :     str	      =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=a				, hidden_states=a    )



class   _lowercase			(  lowerCAmelCase    ):
      _a :  List[Any]										= ResNetConfig
      _a :  Dict										= '''resnet'''
      _a :  int										= '''pixel_values'''


      @property
      def _UpperCamelCase    (     self					:						Tuple    ):





            """simple docstring"""







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



UpperCamelCase_						:					List[str]      				=			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.
"""


UpperCamelCase_						:					Union[str, Any]      				=			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   _lowercase			(  tf.keras.layers.Layer    ):
      _a :  Tuple										= ResNetConfig


      def __init__(     self					:						Tuple				, a					:						ResNetConfig				, **a					:						Any    ):





            """simple docstring"""







            super().__init__(**a    )
            __snake_case   :     Optional[Any]	      =config
            __snake_case   :     Union[str, Any]	      =TFResNetEmbeddings(a				, name='''embedder'''    )
            __snake_case   :     str	      =TFResNetEncoder(a				, name='''encoder'''    )
            __snake_case   :     Optional[int]	      =tf.keras.layers.GlobalAveragePoolingaD(keepdims=a    )







      @unpack_inputs
      def _UpperCamelCase    (     self					:						str				, a					:						tf.Tensor				, a					:						Optional[bool] = None				, a					:						Optional[bool] = None				, a					:						bool = False				, ):





            """simple docstring"""







            __snake_case   :     Optional[Any]	      =(
                output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
            )
            __snake_case   :     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)
            __snake_case   :     Optional[int]	      =tf.transpose(a				, perm=[0, 2, 3, 1]    )
            __snake_case   :     List[Any]	      =self.embedder(a				, training=a    )

            __snake_case   :     int	      =self.encoder(
                a				, output_hidden_states=a				, return_dict=a				, training=a    )

            __snake_case   :     Union[str, Any]	      =encoder_outputs[0]

            __snake_case   :     int	      =self.pooler(a    )

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

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

            __snake_case   :     List[str]	      =hidden_states if output_hidden_states else None

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



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


      def __init__(     self					:						Dict				, a					:						ResNetConfig				, **a					:						int    ):





            """simple docstring"""







            super().__init__(a				, **a    )
            __snake_case   :     Optional[int]	      =TFResNetMainLayer(config=a				, name='''resnet'''    )







      @add_start_docstrings_to_model_forward(a    )
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC				, output_type=a				, config_class=_CONFIG_FOR_DOC				, modality='''vision'''				, expected_output=_EXPECTED_OUTPUT_SHAPE				, )
      @unpack_inputs
      def _UpperCamelCase    (     self					:						Tuple				, a					:						tf.Tensor				, a					:						Optional[bool] = None				, a					:						Optional[bool] = None				, a					:						bool = False				, ):





            """simple docstring"""







            __snake_case   :     Tuple	      =(
                output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
            )
            __snake_case   :     str	      =return_dict if return_dict is not None else self.config.use_return_dict

            __snake_case   :     int	      =self.resnet(
                pixel_values=a				, output_hidden_states=a				, return_dict=a				, training=a				, )
            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.
    '''		,	lowerCAmelCase		,	)
class   _lowercase			(  lowerCAmelCase		,	lowerCAmelCase    ):


      def __init__(     self					:						int				, a					:						ResNetConfig				, **a					:						List[Any]    ):





            """simple docstring"""







            super().__init__(a				, **a    )
            __snake_case   :     Union[str, Any]	      =config.num_labels
            __snake_case   :     int	      =TFResNetMainLayer(a				, name='''resnet'''    )
            # classification head
            __snake_case   :     Optional[int]	      =(
                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 _UpperCamelCase    (     self					:						Union[str, Any]				, a					:						tf.Tensor    ):





            """simple docstring"""







            __snake_case   :     int	      =tf.keras.layers.Flatten()(a    )
            __snake_case   :     int	      =self.classifier_layer(a    )
            return logits







      @add_start_docstrings_to_model_forward(a    )
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT				, output_type=a				, config_class=_CONFIG_FOR_DOC				, expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT				, )
      @unpack_inputs
      def _UpperCamelCase    (     self					:						Dict				, a					:						tf.Tensor = None				, a					:						tf.Tensor = None				, a					:						bool = None				, a					:						bool = None				, a					:						bool = False				, ):





            """simple docstring"""







            __snake_case   :     Optional[int]	      =return_dict if return_dict is not None else self.config.use_return_dict

            __snake_case   :     Union[str, Any]	      =self.resnet(
                a				, output_hidden_states=a				, return_dict=a				, training=a    )

            __snake_case   :     Dict	      =outputs.pooler_output if return_dict else outputs[1]

            __snake_case   :     Union[str, Any]	      =self.classifier(a    )

            __snake_case   :     Dict	      =None if labels is None else self.hf_compute_loss(a				, a    )

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

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










"""simple docstring"""





import numpy as np



def      __lowercase       (     a			:					Union[str, Any]							,		a			:					Optional[Any]							,		a			:					List[str]			)       ->       Dict:
      __snake_case   :     Dict	      =n_draws * class_counts / class_counts.sum()
      # floored means we don't overshoot n_samples, but probably undershoot
      __snake_case   :     int	      =np.floor(a			)
      # we add samples according to how much "left over" probability
      # they had, until we arrive at n_samples
      __snake_case   :     Dict	      =int(n_draws - floored.sum()			)
      if need_to_add > 0:
            __snake_case   :     List[str]	      =continuous - floored
            __snake_case   :     str	      =np.sort(np.unique(a			)			)[::-1]
            # add according to remainder, but break ties
            # randomly to avoid biases
            for value in values:
                  ((__snake_case)      ,					)   :     Tuple	      =np.where(remainder == value			)
                  # if we need_to_add less than what's in inds
                  # we draw randomly from them.
                  # if we need to add more, we add them all and
                  # go to the next value
                  __snake_case   :     int	      =min(len(a			)							,		a			)
                  __snake_case   :     Any	      =rng.choice(a							,		size=a							,		replace=a			)
                  floored[inds] += 1
                  need_to_add -= add_now
                  if need_to_add == 0:
                        break
      return floored.astype(np.intaa			)



def      __lowercase       (     a			:					Any							,		a			:					str							,		a			:					Dict							,		a			:					Union[str, Any]							,		a			:					Union[str, Any]=10			)       ->       Any:
      __snake_case      ,					__snake_case   :     int	      =np.unique(a							,		return_inverse=a			)
      __snake_case   :     Dict	      =classes.shape[0]
      __snake_case   :     int	      =np.bincount(a			)
      if np.min(a			) < 2:
            raise ValueError('''Minimum class count error'''			)
      if n_train < n_classes:
            raise ValueError(
                '''The train_size = %d should be greater or ''' '''equal to the number of classes = %d''' % (n_train, n_classes)			)
      if n_test < n_classes:
            raise ValueError(
                '''The test_size = %d should be greater or ''' '''equal to the number of classes = %d''' % (n_test, n_classes)			)
      __snake_case   :     str	      =np.split(np.argsort(a							,		kind='''mergesort'''			)							,		np.cumsum(a			)[:-1]			)
      for _ in range(a			):
            __snake_case   :     str	      =approximate_mode(a							,		a							,		a			)
            __snake_case   :     List[str]	      =class_counts - n_i
            __snake_case   :     Union[str, Any]	      =approximate_mode(a							,		a							,		a			)

            __snake_case   :     int	      =[]
            __snake_case   :     Union[str, Any]	      =[]

            for i in range(a			):
                  __snake_case   :     Union[str, Any]	      =rng.permutation(class_counts[i]			)
                  __snake_case   :     Dict	      =class_indices[i].take(a							,		mode='''clip'''			)
                  train.extend(perm_indices_class_i[: n_i[i]]			)
                  test.extend(perm_indices_class_i[n_i[i] : n_i[i] + t_i[i]]			)
            __snake_case   :     Any	      =rng.permutation(a			)
            __snake_case   :     int	      =rng.permutation(a			)

            yield train, test


