

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


a_         =      logging.get_logger(__name__)

# General docstring
a_         =      'ResNetConfig'

# Base docstring
a_         =      'microsoft/resnet-50'
a_         =      [1, 2048, 7, 7]

# Image classification docstring
a_         =      'microsoft/resnet-50'
a_         =      'tiger cat'

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






class _UpperCamelCase							(							tf.keras.layers.Layer				):
			'''simple docstring'''







			def __init__(     self							:	Optional[int] , a							:	int , a							:	int = 3 , a							:	int = 1 , a							:	str = "relu" , **a							:	int		)							->    None:

				"""simple docstring"""

				super().__init__(**a		)
				SCREAMING_SNAKE_CASE				: str			       =    kernel_size // 2
				SCREAMING_SNAKE_CASE				: Optional[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
				SCREAMING_SNAKE_CASE				: List[str]			       =    tf.keras.layers.BatchNormalization(epsilon=1e-5 , momentum=0.9 , name="normalization"		)
				SCREAMING_SNAKE_CASE				: Any			       =    ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("linear"		)
			def 							__UpperCamelCase	(     self							:	Dict , a							:	tf.Tensor		)							->    tf.Tensor:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE				: Optional[int]			       =    (self.pad_value, self.pad_value)
				SCREAMING_SNAKE_CASE				: Dict			       =    tf.pad(a , [(0, 0), height_pad, width_pad, (0, 0)]		)
				SCREAMING_SNAKE_CASE				: List[str]			       =    self.conv(a		)
				return hidden_state






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

				"""simple docstring"""

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






class _UpperCamelCase							(							tf.keras.layers.Layer				):
			'''simple docstring'''







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

				"""simple docstring"""

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






			def 							__UpperCamelCase	(     self							:	Tuple , a							:	tf.Tensor , a							:	bool = False		)							->    tf.Tensor:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE				: int			       =    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."		)
				SCREAMING_SNAKE_CASE				: Tuple			       =    pixel_values
				SCREAMING_SNAKE_CASE				: Tuple			       =    self.embedder(a		)
				SCREAMING_SNAKE_CASE				: Dict			       =    tf.pad(a , [[0, 0], [1, 1], [1, 1], [0, 0]]		)
				SCREAMING_SNAKE_CASE				: List[Any]			       =    self.pooler(a		)
				return hidden_state






class _UpperCamelCase							(							tf.keras.layers.Layer				):
			'''simple docstring'''







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

				"""simple docstring"""

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






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

				"""simple docstring"""

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






class _UpperCamelCase							(							tf.keras.layers.Layer				):
			'''simple docstring'''







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

				"""simple docstring"""

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






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

				"""simple docstring"""

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






class _UpperCamelCase							(							tf.keras.layers.Layer				):
			'''simple docstring'''







			def __init__(     self							:	Tuple , a							:	int , a							:	int , a							:	int = 1 , a							:	str = "relu" , a							:	int = 4 , **a							:	Any , )							->    None:

				"""simple docstring"""

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






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

				"""simple docstring"""

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






class _UpperCamelCase							(							tf.keras.layers.Layer				):
			'''simple docstring'''







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

				"""simple docstring"""

				super().__init__(**a		)

				SCREAMING_SNAKE_CASE				: str			       =    TFResNetBottleNeckLayer if config.layer_type == "bottleneck" else TFResNetBasicLayer

				SCREAMING_SNAKE_CASE				: Optional[Any]			       =    [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		)
				]
				SCREAMING_SNAKE_CASE				: int			       =    layers






			def 							__UpperCamelCase	(     self							:	Dict , a							:	tf.Tensor , a							:	bool = False		)							->    tf.Tensor:

				"""simple docstring"""

				for layer in self.stage_layers:
					SCREAMING_SNAKE_CASE				: Dict			       =    layer(a , training=a		)
				return hidden_state






class _UpperCamelCase							(							tf.keras.layers.Layer				):
			'''simple docstring'''







			def __init__(     self							:	List[str] , a							:	ResNetConfig , **a							:	Tuple		)							->    None:

				"""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
				SCREAMING_SNAKE_CASE				: str			       =    [
				    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							:	Dict , a							:	tf.Tensor , a							:	bool = False , a							:	bool = True , a							:	bool = False , )							->    TFBaseModelOutputWithNoAttention:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    () if output_hidden_states else None

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

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

				if output_hidden_states:
					SCREAMING_SNAKE_CASE				: Tuple			       =    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 _UpperCamelCase							(							__A				):
			'''simple docstring'''







			lowerCamelCase__          =ResNetConfig
			lowerCamelCase__          ='resnet'
			lowerCamelCase__          ='pixel_values'
			@property
			def 							__UpperCamelCase	(     self							:	Dict		)							->    str:

				"""simple docstring"""

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

a_         =      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'


a_         =      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 _UpperCamelCase							(							tf.keras.layers.Layer				):
			'''simple docstring'''







			lowerCamelCase__          =ResNetConfig
			def __init__(     self							:	str , a							:	ResNetConfig , **a							:	Dict		)							->    None:

				"""simple docstring"""

				super().__init__(**a		)
				SCREAMING_SNAKE_CASE				: Dict			       =    config
				SCREAMING_SNAKE_CASE				: Dict			       =    TFResNetEmbeddings(a , name="embedder"		)
				SCREAMING_SNAKE_CASE				: Dict			       =    TFResNetEncoder(a , name="encoder"		)
				SCREAMING_SNAKE_CASE				: Dict			       =    tf.keras.layers.GlobalAveragePoolingaD(keepdims=a		)






			@unpack_inputs
			def 							__UpperCamelCase	(     self							:	int , a							:	tf.Tensor , a							:	Optional[bool] = None , a							:	Optional[bool] = None , a							:	bool = False , )							->    Union[Tuple[tf.Tensor], TFBaseModelOutputWithPoolingAndNoAttention]:

				"""simple docstring"""

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

				SCREAMING_SNAKE_CASE				: Dict			       =    self.encoder(
				    a , output_hidden_states=a , return_dict=a , training=a		)

				SCREAMING_SNAKE_CASE				: List[str]			       =    encoder_outputs[0]

				SCREAMING_SNAKE_CASE				: str			       =    self.pooler(a		)

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

				SCREAMING_SNAKE_CASE				: Tuple			       =    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.'							,		__A							,		)
class _UpperCamelCase							(							__A				):
			'''simple docstring'''







			def __init__(     self							:	Any , a							:	ResNetConfig , **a							:	Dict		)							->    None:

				"""simple docstring"""

				super().__init__(a , **a		)
				SCREAMING_SNAKE_CASE				: Optional[Any]			       =    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 , )							->    Union[Tuple[tf.Tensor], TFBaseModelOutputWithPoolingAndNoAttention]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE				: Optional[Any]			       =    (
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				SCREAMING_SNAKE_CASE				: str			       =    return_dict if return_dict is not None else self.config.use_return_dict

				SCREAMING_SNAKE_CASE				: str			       =    self.resnet(
				    pixel_values=a , output_hidden_states=a , return_dict=a , training=a , )
				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    '							,		__A							,		)
class _UpperCamelCase							(							__A							,		__A				):
			'''simple docstring'''







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

				"""simple docstring"""

				super().__init__(a , **a		)
				SCREAMING_SNAKE_CASE				: Any			       =    config.num_labels
				SCREAMING_SNAKE_CASE				: Any			       =    TFResNetMainLayer(a , name="resnet"		)
				# classification head
				SCREAMING_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							:	List[Any] , a							:	tf.Tensor		)							->    tf.Tensor:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE				: Any			       =    tf.keras.layers.Flatten()(a		)
				SCREAMING_SNAKE_CASE				: Dict			       =    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							:	Union[str, Any] , a							:	tf.Tensor = None , a							:	tf.Tensor = None , a							:	bool = None , a							:	bool = None , a							:	bool = False , )							->    Union[Tuple[tf.Tensor], TFImageClassifierOutputWithNoAttention]:

				"""simple docstring"""

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

				SCREAMING_SNAKE_CASE				: str			       =    self.resnet(
				    a , output_hidden_states=a , return_dict=a , training=a		)

				SCREAMING_SNAKE_CASE				: Optional[int]			       =    outputs.pooler_output if return_dict else outputs[1]

				SCREAMING_SNAKE_CASE				: List[Any]			       =    self.classifier(a		)

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

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

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


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


a_         =      {
    'configuration_codegen': ['CODEGEN_PRETRAINED_CONFIG_ARCHIVE_MAP', 'CodeGenConfig', 'CodeGenOnnxConfig'],
    'tokenization_codegen': ['CodeGenTokenizer'],
}

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			a_         =      ['CodeGenTokenizerFast']

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			a_         =      [
			    'CODEGEN_PRETRAINED_MODEL_ARCHIVE_LIST',
			    'CodeGenForCausalLM',
			    'CodeGenModel',
			    'CodeGenPreTrainedModel',
			]

if TYPE_CHECKING:
			from .configuration_codegen import CODEGEN_PRETRAINED_CONFIG_ARCHIVE_MAP, CodeGenConfig, CodeGenOnnxConfig
			from .tokenization_codegen import CodeGenTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_codegen_fast import CodeGenTokenizerFast

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_codegen import (
						    CODEGEN_PRETRAINED_MODEL_ARCHIVE_LIST,
						    CodeGenForCausalLM,
						    CodeGenModel,
						    CodeGenPreTrainedModel,
						)

else:
			import sys

			a_         =      _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)