


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


snake_case_  :     List[Any]							  =							logging.get_logger(__name__)

# General docstring
snake_case_  :     Any							  =							"ResNetConfig"

# Base docstring
snake_case_  :     List[str]							  =							"microsoft/resnet-50"
snake_case_  :     Dict							  =							[1, 20_48, 7, 7]

# Image classification docstring
snake_case_  :     Optional[int]							  =							"microsoft/resnet-50"
snake_case_  :     Dict							  =							"tiger cat"

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


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




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


										"""simple docstring"""





										super().__init__(**UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=kernel_size // 2
										_SCREAMING_SNAKE_CASE     					=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
										_SCREAMING_SNAKE_CASE     					=tf.keras.layers.BatchNormalization(epsilon=1E-5			,      momentum=0.9			,      name='''normalization'''					)
										_SCREAMING_SNAKE_CASE     					=ACTaFN[activation] if activation is not None else tf.keras.layers.Activation('''linear'''					)




						def 							__UpperCamelCase			(    self   : Any			,      _a   : tf.Tensor					)							->       Optional[Any]:


										"""simple docstring"""





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




						def 							__UpperCamelCase			(    self   : Optional[int]			,      _a   : tf.Tensor			,      _a   : bool = False					)							->       int:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=self.convolution(UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=self.normalization(UpperCamelCase_			,      training=UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=self.activation(UpperCamelCase_					)
										return hidden_state


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




						def __init__(    self   : int			,      _a   : ResNetConfig			,      **_a   : Tuple					)							->       Optional[Any]:


										"""simple docstring"""





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




						def 							__UpperCamelCase			(    self   : Any			,      _a   : tf.Tensor			,      _a   : bool = False					)							->       Any:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE					,			_SCREAMING_SNAKE_CASE					,			_SCREAMING_SNAKE_CASE					,			_SCREAMING_SNAKE_CASE     					=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.'''					)
										_SCREAMING_SNAKE_CASE     					=pixel_values
										_SCREAMING_SNAKE_CASE     					=self.embedder(UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=tf.pad(UpperCamelCase_			,      [[0, 0], [1, 1], [1, 1], [0, 0]]					)
										_SCREAMING_SNAKE_CASE     					=self.pooler(UpperCamelCase_					)
										return hidden_state


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




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


										"""simple docstring"""





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




						def 							__UpperCamelCase			(    self   : str			,      _a   : tf.Tensor			,      _a   : bool = False					)							->       int:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=x
										_SCREAMING_SNAKE_CASE     					=self.convolution(UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=self.normalization(UpperCamelCase_			,      training=UpperCamelCase_					)
										return hidden_state


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




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


										"""simple docstring"""





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




						def 							__UpperCamelCase			(    self   : Union[str, Any]			,      _a   : tf.Tensor			,      _a   : bool = False					)							->       Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=hidden_state
										_SCREAMING_SNAKE_CASE     					=self.conva(UpperCamelCase_			,      training=UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=self.conva(UpperCamelCase_			,      training=UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=self.shortcut(UpperCamelCase_			,      training=UpperCamelCase_					)
										hidden_state += residual
										_SCREAMING_SNAKE_CASE     					=self.activation(UpperCamelCase_					)
										return hidden_state


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




						def __init__(    self   : Optional[Any]			,      _a   : int			,      _a   : int			,      _a   : int = 1			,      _a   : str = "relu"			,      _a   : int = 4			,      **_a   : str			,      )							->       Any:


										"""simple docstring"""





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




						def 							__UpperCamelCase			(    self   : Any			,      _a   : tf.Tensor			,      _a   : bool = False					)							->       List[str]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=hidden_state
										_SCREAMING_SNAKE_CASE     					=self.conva(UpperCamelCase_			,      training=UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=self.conva(UpperCamelCase_			,      training=UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=self.conva(UpperCamelCase_			,      training=UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=self.shortcut(UpperCamelCase_			,      training=UpperCamelCase_					)
										hidden_state += residual
										_SCREAMING_SNAKE_CASE     					=self.activation(UpperCamelCase_					)
										return hidden_state


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




						def __init__(    self   : str			,      _a   : ResNetConfig			,      _a   : int			,      _a   : int			,      _a   : int = 2			,      _a   : int = 2			,      **_a   : Optional[int]					)							->       List[Any]:


										"""simple docstring"""





										super().__init__(**UpperCamelCase_					)

										_SCREAMING_SNAKE_CASE     					=TFResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else TFResNetBasicLayer

										_SCREAMING_SNAKE_CASE     					=[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					)
										]
										_SCREAMING_SNAKE_CASE     					=layers




						def 							__UpperCamelCase			(    self   : List[str]			,      _a   : tf.Tensor			,      _a   : bool = False					)							->       str:


										"""simple docstring"""





										for layer in self.stage_layers:
														_SCREAMING_SNAKE_CASE     					=layer(UpperCamelCase_			,      training=UpperCamelCase_					)
										return hidden_state


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




						def __init__(    self   : int			,      _a   : ResNetConfig			,      **_a   : Dict					)							->       Dict:


										"""simple docstring"""





										super().__init__(**UpperCamelCase_					)
										# based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
										_SCREAMING_SNAKE_CASE     					=[
										    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 							__UpperCamelCase			(    self   : Optional[Any]			,      _a   : tf.Tensor			,      _a   : bool = False			,      _a   : bool = True			,      _a   : bool = False			,      )							->       str:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=() if output_hidden_states else None

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

														_SCREAMING_SNAKE_CASE     					=stage_module(UpperCamelCase_			,      training=UpperCamelCase_					)

										if output_hidden_states:
														_SCREAMING_SNAKE_CASE     					=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__							(						lowercase_					):
						UpperCAmelCase     =  ResNetConfig
						UpperCAmelCase     =  "resnet"
						UpperCAmelCase     =  "pixel_values"




						@property
						def 							__UpperCamelCase			(    self   : str					)							->       Tuple:


										"""simple docstring"""





										return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)			,      dtype=tf.floataa					)}
snake_case_  :     Union[str, Any]							  =							R"\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\n    Parameters:\n        config ([`ResNetConfig`]): 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"


snake_case_  :     Tuple							  =							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            [`ConvNextImageProcessor.__call__`] for details.\n\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"


@keras_serializable
class   A__							(						tf.keras.layers.Layer					):
						UpperCAmelCase     =  ResNetConfig




						def __init__(    self   : str			,      _a   : ResNetConfig			,      **_a   : str					)							->       Dict:


										"""simple docstring"""





										super().__init__(**UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=config
										_SCREAMING_SNAKE_CASE     					=TFResNetEmbeddings(UpperCamelCase_			,      name='''embedder'''					)
										_SCREAMING_SNAKE_CASE     					=TFResNetEncoder(UpperCamelCase_			,      name='''encoder'''					)
										_SCREAMING_SNAKE_CASE     					=tf.keras.layers.GlobalAveragePoolingaD(keepdims=UpperCamelCase_					)




						@unpack_inputs
						def 							__UpperCamelCase			(    self   : str			,      _a   : tf.Tensor			,      _a   : Optional[bool] = None			,      _a   : Optional[bool] = None			,      _a   : bool = False			,      )							->       Optional[int]:


										"""simple docstring"""





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

										_SCREAMING_SNAKE_CASE     					=self.encoder(
										    UpperCamelCase_			,      output_hidden_states=UpperCamelCase_			,      return_dict=UpperCamelCase_			,      training=UpperCamelCase_					)

										_SCREAMING_SNAKE_CASE     					=encoder_outputs[0]

										_SCREAMING_SNAKE_CASE     					=self.pooler(UpperCamelCase_					)

										# Transpose all the outputs to the NCHW format
										# (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
										_SCREAMING_SNAKE_CASE     					=tf.transpose(UpperCamelCase_			,      (0, 3, 1, 2)					)
										_SCREAMING_SNAKE_CASE     					=tf.transpose(UpperCamelCase_			,      (0, 3, 1, 2)					)
										_SCREAMING_SNAKE_CASE     					=()
										for hidden_state in encoder_outputs[1:]:
														_SCREAMING_SNAKE_CASE     					=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

										_SCREAMING_SNAKE_CASE     					=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."				,       lowercase_				,       )
class   A__							(						lowercase_					):




						def __init__(    self   : int			,      _a   : ResNetConfig			,      **_a   : Dict					)							->       Union[str, Any]:


										"""simple docstring"""





										super().__init__(UpperCamelCase_			,      **UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=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 							__UpperCamelCase			(    self   : Dict			,      _a   : tf.Tensor			,      _a   : Optional[bool] = None			,      _a   : Optional[bool] = None			,      _a   : bool = False			,      )							->       Dict:


										"""simple docstring"""





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

										_SCREAMING_SNAKE_CASE     					=self.resnet(
										    pixel_values=UpperCamelCase_			,      output_hidden_states=UpperCamelCase_			,      return_dict=UpperCamelCase_			,      training=UpperCamelCase_			,      )
										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   A__							(						lowercase_				,       lowercase_					):




						def __init__(    self   : int			,      _a   : ResNetConfig			,      **_a   : Any					)							->       int:


										"""simple docstring"""





										super().__init__(UpperCamelCase_			,      **UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=config.num_labels
										_SCREAMING_SNAKE_CASE     					=TFResNetMainLayer(UpperCamelCase_			,      name='''resnet'''					)
										# classification head
										_SCREAMING_SNAKE_CASE     					=(
										    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					)							->       List[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=tf.keras.layers.Flatten()(UpperCamelCase_					)
										_SCREAMING_SNAKE_CASE     					=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 							__UpperCamelCase			(    self   : str			,      _a   : tf.Tensor = None			,      _a   : tf.Tensor = None			,      _a   : bool = None			,      _a   : bool = None			,      _a   : bool = False			,      )							->       int:


										"""simple docstring"""





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

										_SCREAMING_SNAKE_CASE     					=self.resnet(
										    UpperCamelCase_			,      output_hidden_states=UpperCamelCase_			,      return_dict=UpperCamelCase_			,      training=UpperCamelCase_					)

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

										_SCREAMING_SNAKE_CASE     					=self.classifier(UpperCamelCase_					)

										_SCREAMING_SNAKE_CASE     					=None if labels is None else self.hf_compute_loss(UpperCamelCase_			,      UpperCamelCase_					)

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

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



import math

class   A__							:




						def __init__(    self   : Optional[int]			,      _a   : int			,      _a   : List[Any]			,      _a   : int					)							->       Optional[int]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=x
										_SCREAMING_SNAKE_CASE     					=y
										_SCREAMING_SNAKE_CASE     					=z




						def __repr__(    self   : List[str]					)							->       str:


										"""simple docstring"""





										return f"Point({self.x}, {self.y}, {self.z})"







def   lowerCamelCase(				a__		,a__):
				return math.sqrt(abs((b.x - a.x) ** 2 + (b.y - a.y) ** 2 + (b.z - a.z) ** 2))


if __name__ == "__main__":
							import doctest

							doctest.testmod()