

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


lowerCAmelCase__             =		logging.get_logger(__name__)

# General docstring
lowerCAmelCase__             =		"""ResNetConfig"""

# Base docstring
lowerCAmelCase__             =		"""microsoft/resnet-50"""
lowerCAmelCase__             =		[1, 2_0_4_8, 7, 7]

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

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







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


		"""simple docstring"""




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


		def 		A__       (					self						,							__snake_case):
			# Pad to match that done in the PyTorch Conv2D model
			_UpperCamelCase :		Optional[Any]								=    (self.pad_value, self.pad_value)
			_UpperCamelCase :		Union[str, Any]								=    tf.pad(__snake_case						,							[(0, 0), height_pad, width_pad, (0, 0)])
			_UpperCamelCase :		int								=    self.conv(__snake_case)
			return hidden_state






		def 		A__       (					self						,							__snake_case						,							__snake_case = False):
			_UpperCamelCase :		Optional[Any]								=    self.convolution(__snake_case)
			_UpperCamelCase :		str								=    self.normalization(__snake_case						,							training=__snake_case)
			_UpperCamelCase :		int								=    self.activation(__snake_case)
			return hidden_state







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


		"""simple docstring"""




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






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







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


		"""simple docstring"""




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






		def 		A__       (					self						,							__snake_case						,							__snake_case = False):
			_UpperCamelCase :		Dict								=    x
			_UpperCamelCase :		Union[str, Any]								=    self.convolution(__snake_case)
			_UpperCamelCase :		List[Any]								=    self.normalization(__snake_case						,							training=__snake_case)
			return hidden_state







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


		"""simple docstring"""




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






		def 		A__       (					self						,							__snake_case						,							__snake_case = False):
			_UpperCamelCase :		Union[str, Any]								=    hidden_state
			_UpperCamelCase :		Optional[Any]								=    self.conva(__snake_case						,							training=__snake_case)
			_UpperCamelCase :		List[str]								=    self.conva(__snake_case						,							training=__snake_case)
			_UpperCamelCase :		int								=    self.shortcut(__snake_case						,							training=__snake_case)
			hidden_state += residual
			_UpperCamelCase :		Tuple								=    self.activation(__snake_case)
			return hidden_state







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


		"""simple docstring"""




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






		def 		A__       (					self						,							__snake_case						,							__snake_case = False):
			_UpperCamelCase :		Any								=    hidden_state
			_UpperCamelCase :		List[Any]								=    self.conva(__snake_case						,							training=__snake_case)
			_UpperCamelCase :		int								=    self.conva(__snake_case						,							training=__snake_case)
			_UpperCamelCase :		int								=    self.conva(__snake_case						,							training=__snake_case)
			_UpperCamelCase :		Tuple								=    self.shortcut(__snake_case						,							training=__snake_case)
			hidden_state += residual
			_UpperCamelCase :		Optional[Any]								=    self.activation(__snake_case)
			return hidden_state







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


		"""simple docstring"""




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

			_UpperCamelCase :		List[Any]								=    TFResNetBottleNeckLayer if config.layer_type == 'bottleneck' else TFResNetBasicLayer

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






		def 		A__       (					self						,							__snake_case						,							__snake_case = False):
			for layer in self.stage_layers:
				_UpperCamelCase :		Optional[int]								=    layer(__snake_case						,							training=__snake_case)
			return hidden_state







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


		"""simple docstring"""




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






		def 		A__       (					self						,							__snake_case						,							__snake_case = False						,							__snake_case = True						,							__snake_case = False						,							):
			_UpperCamelCase :		List[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 :		Dict								=    stage_module(__snake_case						,							training=__snake_case)

			if output_hidden_states:
				_UpperCamelCase :		Optional[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)







class 						lowercase     (					_lowercase				):


		"""simple docstring"""




		a__       							=	ResNetConfig
		a__       							=	"resnet"
		a__       							=	"pixel_values"


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

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


lowerCAmelCase__             =		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				):


		"""simple docstring"""




		a__       							=	ResNetConfig


		def __init__(					self						,							__snake_case						,							**__snake_case):
			super().__init__(**__snake_case)
			_UpperCamelCase :		List[Any]								=    config
			_UpperCamelCase :		Optional[Any]								=    TFResNetEmbeddings(__snake_case						,							name='embedder')
			_UpperCamelCase :		Any								=    TFResNetEncoder(__snake_case						,							name='encoder')
			_UpperCamelCase :		str								=    tf.keras.layers.GlobalAveragePoolingaD(keepdims=__snake_case)






		@unpack_inputs
		def 		A__       (					self						,							__snake_case						,							__snake_case = None						,							__snake_case = None						,							__snake_case = False						,							):
			_UpperCamelCase :		Tuple								=    (
			    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

			# 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 :		Dict								=    tf.transpose(__snake_case						,							perm=[0, 2, 3, 1])
			_UpperCamelCase :		str								=    self.embedder(__snake_case						,							training=__snake_case)

			_UpperCamelCase :		Any								=    self.encoder(
			    __snake_case						,							output_hidden_states=__snake_case						,							return_dict=__snake_case						,							training=__snake_case)

			_UpperCamelCase :		Tuple								=    encoder_outputs[0]

			_UpperCamelCase :		Optional[Any]								=    self.pooler(__snake_case)

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

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

			_UpperCamelCase :		Tuple								=    hidden_states if output_hidden_states else None

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







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


		"""simple docstring"""




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






		@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						,							)
		@unpack_inputs
		def 		A__       (					self						,							__snake_case						,							__snake_case = None						,							__snake_case = None						,							__snake_case = False						,							):
			_UpperCamelCase :		Any								=    (
			    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
			)
			_UpperCamelCase :		Optional[int]								=    return_dict if return_dict is not None else self.config.use_return_dict

			_UpperCamelCase :		int								=    self.resnet(
			    pixel_values=__snake_case						,							output_hidden_states=__snake_case						,							return_dict=__snake_case						,							training=__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    "     ,				_lowercase     ,				)
class 						lowercase     (					_lowercase     ,				_lowercase				):


		"""simple docstring"""




		def __init__(					self						,							__snake_case						,							**__snake_case):
			super().__init__(__snake_case						,							**__snake_case)
			_UpperCamelCase :		List[Any]								=    config.num_labels
			_UpperCamelCase :		List[str]								=    TFResNetMainLayer(__snake_case						,							name='resnet')
			# classification head
			_UpperCamelCase :		Tuple								=    (
			    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 		A__       (					self						,							__snake_case):
			_UpperCamelCase :		Tuple								=    tf.keras.layers.Flatten()(__snake_case)
			_UpperCamelCase :		Dict								=    self.classifier_layer(__snake_case)
			return logits






		@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						,							)
		@unpack_inputs
		def 		A__       (					self						,							__snake_case = None						,							__snake_case = None						,							__snake_case = None						,							__snake_case = None						,							__snake_case = False						,							):
			_UpperCamelCase :		List[Any]								=    return_dict if return_dict is not None else self.config.use_return_dict

			_UpperCamelCase :		List[Any]								=    self.resnet(
			    __snake_case						,							output_hidden_states=__snake_case						,							return_dict=__snake_case						,							training=__snake_case)

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

			_UpperCamelCase :		List[str]								=    self.classifier(__snake_case)

			_UpperCamelCase :		Union[str, Any]								=    None if labels is None else self.hf_compute_loss(__snake_case						,							__snake_case)

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

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








"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
    is_vision_available,
)


lowercase  		=			{
    """configuration_mobilevitv2""": [
        """MOBILEVITV2_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """MobileViTV2Config""",
        """MobileViTV2OnnxConfig""",
    ],
}


try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase  		=			[
					    """MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """MobileViTV2ForImageClassification""",
					    """MobileViTV2ForSemanticSegmentation""",
					    """MobileViTV2Model""",
					    """MobileViTV2PreTrainedModel""",
					]

if TYPE_CHECKING:
					from .configuration_mobilevitva import (
					    MOBILEVITV2_PRETRAINED_CONFIG_ARCHIVE_MAP,
					    MobileViTVaConfig,
					    MobileViTVaOnnxConfig,
					)

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_mobilevitva import (
										    MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
										    MobileViTVaForImageClassification,
										    MobileViTVaForSemanticSegmentation,
										    MobileViTVaModel,
										    MobileViTVaPreTrainedModel,
										)

else:
					import sys

					lowercase  		=			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





