






from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


__UpperCamelCase      :      Union[str, Any]     						=						logging.get_logger(__name__)

# General docstring
__UpperCamelCase      :      List[Any]     						=						"""RegNetConfig"""

# Base docstring
__UpperCamelCase      :      Dict     						=						"""facebook/regnet-y-040"""
__UpperCamelCase      :      Tuple     						=						[1, 1088, 7, 7]

# Image classification docstring
__UpperCamelCase      :      Union[str, Any]     						=						"""facebook/regnet-y-040"""
__UpperCamelCase      :      Tuple     						=						"""tabby, tabby cat"""

__UpperCamelCase      :      Optional[int]     						=						[
    """facebook/regnet-y-040""",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]


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







	def __init__(  self      :							str       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							int = 3       ,				_lowerCAmelCase      :							int = 1       ,				_lowerCAmelCase      :							int = 1       ,				_lowerCAmelCase      :							Optional[str] = "relu"       ,				**_lowerCAmelCase      :							Tuple       ,				)		-> str:
		"""simple docstring"""







		super().__init__(**_lowerCAmelCase			)
		# The padding and conv has been verified in
		# https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
		__lowercase             =  tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2			)
		__lowercase             =  tf.keras.layers.ConvaD(
		    filters=_lowerCAmelCase       ,				kernel_size=_lowerCAmelCase       ,				strides=_lowerCAmelCase       ,				padding="""VALID"""       ,				groups=_lowerCAmelCase       ,				use_bias=_lowerCAmelCase       ,				name="""convolution"""       ,				)
		__lowercase             =  tf.keras.layers.BatchNormalization(epsilon=1e-5       ,				momentum=0.9       ,				name="""normalization"""			)
		__lowercase             =  ACTaFN[activation] if activation is not None else tf.identity







	def  _a   (  self      :							int       ,				_lowerCAmelCase      :							Any			)		-> str:
		"""simple docstring"""







		__lowercase             =  self.convolution(self.padding(_lowerCAmelCase			)			)
		__lowercase             =  self.normalization(_lowerCAmelCase			)
		__lowercase             =  self.activation(_lowerCAmelCase			)
		return hidden_state


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







	def __init__(  self      :							int       ,				_lowerCAmelCase      :							RegNetConfig       ,				**_lowerCAmelCase      :							str			)		-> Union[str, Any]:
		"""simple docstring"""







		super().__init__(**_lowerCAmelCase			)
		__lowercase             =  config.num_channels
		__lowercase             =  TFRegNetConvLayer(
		    out_channels=config.embedding_size       ,				kernel_size=3       ,				stride=2       ,				activation=config.hidden_act       ,				name="""embedder"""       ,				)







	def  _a   (  self      :							Optional[int]       ,				_lowerCAmelCase      :							Optional[int]			)		-> Optional[int]:
		"""simple docstring"""







		__lowercase             =  shape_list(_lowerCAmelCase			)[1]
		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."""			)

		# When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
		# So change the input format from `NCHW` to `NHWC`.
		# shape = (batch_size, in_height, in_width, in_channels=num_channels)
		__lowercase             =  tf.transpose(_lowerCAmelCase       ,				perm=(0, 2, 3, 1)			)
		__lowercase             =  self.embedder(_lowerCAmelCase			)
		return hidden_state


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







	def __init__(  self      :							Tuple       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							int = 2       ,				**_lowerCAmelCase      :							Optional[int]			)		-> List[str]:
		"""simple docstring"""







		super().__init__(**_lowerCAmelCase			)
		__lowercase             =  tf.keras.layers.ConvaD(
		    filters=_lowerCAmelCase       ,				kernel_size=1       ,				strides=_lowerCAmelCase       ,				use_bias=_lowerCAmelCase       ,				name="""convolution"""			)
		__lowercase             =  tf.keras.layers.BatchNormalization(epsilon=1e-5       ,				momentum=0.9       ,				name="""normalization"""			)







	def  _a   (  self      :							List[Any]       ,				_lowerCAmelCase      :							tf.Tensor       ,				_lowerCAmelCase      :							bool = False			)		-> tf.Tensor:
		"""simple docstring"""







		return self.normalization(self.convolution(_lowerCAmelCase			)       ,				training=_lowerCAmelCase			)


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







	def __init__(  self      :							Union[str, Any]       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							int       ,				**_lowerCAmelCase      :							Tuple			)		-> Dict:
		"""simple docstring"""







		super().__init__(**_lowerCAmelCase			)
		__lowercase             =  tf.keras.layers.GlobalAveragePoolingaD(keepdims=_lowerCAmelCase       ,				name="""pooler"""			)
		__lowercase             =  [
		    tf.keras.layers.ConvaD(filters=_lowerCAmelCase       ,				kernel_size=1       ,				activation="""relu"""       ,				name="""attention.0"""			),
		    tf.keras.layers.ConvaD(filters=_lowerCAmelCase       ,				kernel_size=1       ,				activation="""sigmoid"""       ,				name="""attention.2"""			),
		]







	def  _a   (  self      :							int       ,				_lowerCAmelCase      :							Tuple			)		-> Optional[Any]:
		"""simple docstring"""







		__lowercase             =  self.pooler(_lowerCAmelCase			)
		for layer_module in self.attention:
			__lowercase             =  layer_module(_lowerCAmelCase			)
		__lowercase             =  hidden_state * pooled
		return hidden_state


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







	def __init__(  self      :							Optional[Any]       ,				_lowerCAmelCase      :							RegNetConfig       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							int = 1       ,				**_lowerCAmelCase      :							Dict			)		-> List[str]:
		"""simple docstring"""







		super().__init__(**_lowerCAmelCase			)
		__lowercase             =  in_channels != out_channels or stride != 1
		__lowercase             =  max(1       ,				out_channels // config.groups_width			)
		__lowercase             =  (
		    TFRegNetShortCut(_lowerCAmelCase       ,				stride=_lowerCAmelCase       ,				name="""shortcut"""			)
		    if should_apply_shortcut
		    else tf.keras.layers.Activation("""linear"""       ,				name="""shortcut"""			)
		)
		# `self.layers` instead of `self.layer` because that is a reserved argument.
		__lowercase             =  [
		    TFRegNetConvLayer(_lowerCAmelCase       ,				kernel_size=1       ,				activation=config.hidden_act       ,				name="""layer.0"""			),
		    TFRegNetConvLayer(
		        _lowerCAmelCase       ,				stride=_lowerCAmelCase       ,				groups=_lowerCAmelCase       ,				activation=config.hidden_act       ,				name="""layer.1"""			),
		    TFRegNetConvLayer(_lowerCAmelCase       ,				kernel_size=1       ,				activation=_lowerCAmelCase       ,				name="""layer.2"""			),
		]
		__lowercase             =  ACTaFN[config.hidden_act]







	def  _a   (  self      :							int       ,				_lowerCAmelCase      :							Tuple			)		-> int:
		"""simple docstring"""







		__lowercase             =  hidden_state
		for layer_module in self.layers:
			__lowercase             =  layer_module(_lowerCAmelCase			)
		__lowercase             =  self.shortcut(_lowerCAmelCase			)
		hidden_state += residual
		__lowercase             =  self.activation(_lowerCAmelCase			)
		return hidden_state


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







	def __init__(  self      :							str       ,				_lowerCAmelCase      :							RegNetConfig       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							int = 1       ,				**_lowerCAmelCase      :							Tuple			)		-> Tuple:
		"""simple docstring"""







		super().__init__(**_lowerCAmelCase			)
		__lowercase             =  in_channels != out_channels or stride != 1
		__lowercase             =  max(1       ,				out_channels // config.groups_width			)
		__lowercase             =  (
		    TFRegNetShortCut(_lowerCAmelCase       ,				stride=_lowerCAmelCase       ,				name="""shortcut"""			)
		    if should_apply_shortcut
		    else tf.keras.layers.Activation("""linear"""       ,				name="""shortcut"""			)
		)
		__lowercase             =  [
		    TFRegNetConvLayer(_lowerCAmelCase       ,				kernel_size=1       ,				activation=config.hidden_act       ,				name="""layer.0"""			),
		    TFRegNetConvLayer(
		        _lowerCAmelCase       ,				stride=_lowerCAmelCase       ,				groups=_lowerCAmelCase       ,				activation=config.hidden_act       ,				name="""layer.1"""			),
		    TFRegNetSELayer(_lowerCAmelCase       ,				reduced_channels=int(round(in_channels / 4			)			)       ,				name="""layer.2"""			),
		    TFRegNetConvLayer(_lowerCAmelCase       ,				kernel_size=1       ,				activation=_lowerCAmelCase       ,				name="""layer.3"""			),
		]
		__lowercase             =  ACTaFN[config.hidden_act]







	def  _a   (  self      :							List[str]       ,				_lowerCAmelCase      :							Optional[int]			)		-> Dict:
		"""simple docstring"""







		__lowercase             =  hidden_state
		for layer_module in self.layers:
			__lowercase             =  layer_module(_lowerCAmelCase			)
		__lowercase             =  self.shortcut(_lowerCAmelCase			)
		hidden_state += residual
		__lowercase             =  self.activation(_lowerCAmelCase			)
		return hidden_state


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







	def __init__(  self      :							Optional[int]       ,				_lowerCAmelCase      :							RegNetConfig       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							int = 2       ,				_lowerCAmelCase      :							int = 2       ,				**_lowerCAmelCase      :							Optional[Any]			)		-> Optional[Any]:
		"""simple docstring"""







		super().__init__(**_lowerCAmelCase			)

		__lowercase             =  TFRegNetXLayer if config.layer_type == """x""" else TFRegNetYLayer
		__lowercase             =  [
		    # downsampling is done in the first layer with stride of 2
		    layer(_lowerCAmelCase       ,				_lowerCAmelCase       ,				_lowerCAmelCase       ,				stride=_lowerCAmelCase       ,				name="""layers.0"""			),
		    *[layer(_lowerCAmelCase       ,				_lowerCAmelCase       ,				_lowerCAmelCase       ,				name=F'layers.{i+1}'			) for i in range(depth - 1			)],
		]







	def  _a   (  self      :							Optional[Any]       ,				_lowerCAmelCase      :							Any			)		-> Optional[int]:
		"""simple docstring"""







		for layer_module in self.layers:
			__lowercase             =  layer_module(_lowerCAmelCase			)
		return hidden_state


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







	def __init__(  self      :							List[Any]       ,				_lowerCAmelCase      :							RegNetConfig       ,				**_lowerCAmelCase      :							Any			)		-> Union[str, Any]:
		"""simple docstring"""







		super().__init__(**_lowerCAmelCase			)
		__lowercase             =  []
		# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
		self.stages.append(
		    TFRegNetStage(
		        _lowerCAmelCase       ,				config.embedding_size       ,				config.hidden_sizes[0]       ,				stride=2 if config.downsample_in_first_stage else 1       ,				depth=config.depths[0]       ,				name="""stages.0"""       ,				)			)
		__lowercase             =  zip(config.hidden_sizes       ,				config.hidden_sizes[1:]			)
		for i, ((in_channels, out_channels), depth) in enumerate(zip(_lowerCAmelCase       ,				config.depths[1:]			)			):
			self.stages.append(TFRegNetStage(_lowerCAmelCase       ,				_lowerCAmelCase       ,				_lowerCAmelCase       ,				depth=_lowerCAmelCase       ,				name=F'stages.{i+1}'			)			)







	def  _a   (  self      :							List[Any]       ,				_lowerCAmelCase      :							tf.Tensor       ,				_lowerCAmelCase      :							bool = False       ,				_lowerCAmelCase      :							bool = True			)		-> TFBaseModelOutputWithNoAttention:
		"""simple docstring"""







		__lowercase             =  () if output_hidden_states else None

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

			__lowercase             =  stage_module(_lowerCAmelCase			)

		if output_hidden_states:
			__lowercase             =  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=_lowerCAmelCase       ,				hidden_states=_lowerCAmelCase			)


@keras_serializable
class        __UpperCamelCase    (     tf.keras.layers.Layer   ):
	__snake_case	:Dict   =       RegNetConfig







	def __init__(  self      :							List[str]       ,				_lowerCAmelCase      :							Optional[int]       ,				**_lowerCAmelCase      :							int			)		-> Union[str, Any]:
		"""simple docstring"""







		super().__init__(**_lowerCAmelCase			)
		__lowercase             =  config
		__lowercase             =  TFRegNetEmbeddings(_lowerCAmelCase       ,				name="""embedder"""			)
		__lowercase             =  TFRegNetEncoder(_lowerCAmelCase       ,				name="""encoder"""			)
		__lowercase             =  tf.keras.layers.GlobalAveragePoolingaD(keepdims=_lowerCAmelCase       ,				name="""pooler"""			)







	@unpack_inputs
	def  _a   (  self      :							List[str]       ,				_lowerCAmelCase      :							tf.Tensor       ,				_lowerCAmelCase      :							Optional[bool] = None       ,				_lowerCAmelCase      :							Optional[bool] = None       ,				_lowerCAmelCase      :							bool = False       ,				)		-> TFBaseModelOutputWithPoolingAndNoAttention:
		"""simple docstring"""







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

		__lowercase             =  self.embedder(_lowerCAmelCase       ,				training=_lowerCAmelCase			)

		__lowercase             =  self.encoder(
		    _lowerCAmelCase       ,				output_hidden_states=_lowerCAmelCase       ,				return_dict=_lowerCAmelCase       ,				training=_lowerCAmelCase			)

		__lowercase             =  encoder_outputs[0]
		__lowercase             =  self.pooler(_lowerCAmelCase			)

		# Change to NCHW output format have uniformity in the modules
		__lowercase             =  tf.transpose(_lowerCAmelCase       ,				perm=(0, 3, 1, 2)			)
		__lowercase             =  tf.transpose(_lowerCAmelCase       ,				perm=(0, 3, 1, 2)			)

		# Change the other hidden state outputs to NCHW as well
		if output_hidden_states:
			__lowercase             =  tuple([tf.transpose(_lowerCAmelCase       ,				perm=(0, 3, 1, 2)			) for h in encoder_outputs[1]]			)

		if not return_dict:
			return (last_hidden_state, pooled_output) + encoder_outputs[1:]

		return TFBaseModelOutputWithPoolingAndNoAttention(
		    last_hidden_state=_lowerCAmelCase       ,				pooler_output=_lowerCAmelCase       ,				hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states       ,				)


class        __UpperCamelCase    (     _lowerCAmelCase   ):
	__snake_case	:Optional[int]   =       RegNetConfig
	__snake_case	:Optional[Any]   =       'regnet'
	__snake_case	:Dict   =       'pixel_values'







	@property
	def  _a   (  self      :							List[str]			)		-> Any:
		"""simple docstring"""







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

__UpperCamelCase      :      Tuple     						=						r"""
    Parameters:
    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.
        config ([`RegNetConfig`]): 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
            [`ConveNextImageProcessor.__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.
"""


@add_start_docstrings(
    'The bare RegNet model outputting raw features without any specific head on top.'   ,			_lowerCAmelCase   ,			)
class        __UpperCamelCase    (     _lowerCAmelCase   ):







	def __init__(  self      :							Optional[Any]       ,				_lowerCAmelCase      :							RegNetConfig       ,				*_lowerCAmelCase      :							Tuple       ,				**_lowerCAmelCase      :							Any			)		-> List[Any]:
		"""simple docstring"""







		super().__init__(_lowerCAmelCase       ,				*_lowerCAmelCase       ,				**_lowerCAmelCase			)
		__lowercase             =  TFRegNetMainLayer(_lowerCAmelCase       ,				name="""regnet"""			)







	@unpack_inputs
	@add_start_docstrings_to_model_forward(_lowerCAmelCase			)
	@add_code_sample_docstrings(
	    checkpoint=_CHECKPOINT_FOR_DOC       ,				output_type=_lowerCAmelCase       ,				config_class=_CONFIG_FOR_DOC       ,				modality="""vision"""       ,				expected_output=_EXPECTED_OUTPUT_SHAPE       ,				)
	def  _a   (  self      :							Tuple       ,				_lowerCAmelCase      :							tf.Tensor       ,				_lowerCAmelCase      :							Optional[bool] = None       ,				_lowerCAmelCase      :							Optional[bool] = None       ,				_lowerCAmelCase      :							int=False       ,				)		-> Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:
		"""simple docstring"""







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

		__lowercase             =  self.regnet(
		    pixel_values=_lowerCAmelCase       ,				output_hidden_states=_lowerCAmelCase       ,				return_dict=_lowerCAmelCase       ,				training=_lowerCAmelCase       ,				)
		if not return_dict:
			return (outputs[0],) + outputs[1:]

		return TFBaseModelOutputWithPoolingAndNoAttention(
		    last_hidden_state=outputs.last_hidden_state       ,				pooler_output=outputs.pooler_output       ,				hidden_states=outputs.hidden_states       ,				)


@add_start_docstrings(
    '\n    RegNet 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      :							int       ,				_lowerCAmelCase      :							RegNetConfig       ,				*_lowerCAmelCase      :							Union[str, Any]       ,				**_lowerCAmelCase      :							Optional[int]			)		-> int:
		"""simple docstring"""







		super().__init__(_lowerCAmelCase       ,				*_lowerCAmelCase       ,				**_lowerCAmelCase			)
		__lowercase             =  config.num_labels
		__lowercase             =  TFRegNetMainLayer(_lowerCAmelCase       ,				name="""regnet"""			)
		# classification head
		__lowercase             =  [
		    tf.keras.layers.Flatten(),
		    tf.keras.layers.Dense(config.num_labels       ,				name="""classifier.1"""			) if config.num_labels > 0 else tf.identity,
		]







	@unpack_inputs
	@add_start_docstrings_to_model_forward(_lowerCAmelCase			)
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT       ,				output_type=_lowerCAmelCase       ,				config_class=_CONFIG_FOR_DOC       ,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT       ,				)
	def  _a   (  self      :							List[Any]       ,				_lowerCAmelCase      :							tf.Tensor = None       ,				_lowerCAmelCase      :							tf.Tensor = None       ,				_lowerCAmelCase      :							bool = None       ,				_lowerCAmelCase      :							bool = None       ,				_lowerCAmelCase      :							Union[str, Any]=False       ,				)		-> Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
		"""simple docstring"""







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

		__lowercase             =  self.regnet(
		    _lowerCAmelCase       ,				output_hidden_states=_lowerCAmelCase       ,				return_dict=_lowerCAmelCase       ,				training=_lowerCAmelCase			)

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

		__lowercase             =  self.classifier[0](_lowerCAmelCase			)
		__lowercase             =  self.classifier[1](_lowerCAmelCase			)

		__lowercase             =  None if labels is None else self.hf_compute_loss(labels=_lowerCAmelCase       ,				logits=_lowerCAmelCase			)

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

		return TFSequenceClassifierOutput(loss=_lowerCAmelCase       ,				logits=_lowerCAmelCase       ,				hidden_states=outputs.hidden_states			)









from __future__ import annotations





def 				snake_case  (     lowerCamelCase					):




	'''simple docstring'''







	if not nums:
		return 0
	__lowercase             =  nums[0]
	__lowercase             =  0
	for num in nums[1:]:
		__lowercase       ,						__lowercase             =  (
		    max_excluding + num,
		    max(lowerCamelCase				,  lowerCamelCase					),
		)
	return max(lowerCamelCase				,  lowerCamelCase					)


if __name__ == "__main__":
						import doctest

						doctest.testmod()

