


"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A     :   str			    =			{
    "configuration_pegasus_x": ["PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP", "PegasusXConfig"],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	A     :   Dict			    =			[
	    "PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST",
	    "PegasusXForConditionalGeneration",
	    "PegasusXModel",
	    "PegasusXPreTrainedModel",
	]


if TYPE_CHECKING:
	from .configuration_pegasus_x import PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusXConfig

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_pegasus_x import (
		    PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST,
		    PegasusXForConditionalGeneration,
		    PegasusXModel,
		    PegasusXPreTrainedModel,
		)


else:
	import sys

	A     :   Any			    =			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






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


__a  			=					logging.get_logger(__name__)

# General docstring
__a  			=					'RegNetConfig'

# Base docstring
__a  			=					'facebook/regnet-y-040'
__a  			=					[1, 1_0_8_8, 7, 7]

# Image classification docstring
__a  			=					'facebook/regnet-y-040'
__a  			=					'tabby, tabby cat'

__a  			=					[
    'facebook/regnet-y-040',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]
class  lowercase__(				tf.keras.layers.Layer							):



	"""simple docstring"""




	def __init__(       self  :			Union[str, Any]       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int = 3       ,							SCREAMING_SNAKE_CASE_  :			int = 1       ,							SCREAMING_SNAKE_CASE_  :			int = 1       ,							SCREAMING_SNAKE_CASE_  :			Optional[str] = "relu"       ,							**SCREAMING_SNAKE_CASE_  :			Any       ,							)       ->     List[str]:
					super().__init__(**SCREAMING_SNAKE_CASE_       )
					# 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=SCREAMING_SNAKE_CASE_       ,							kernel_size=SCREAMING_SNAKE_CASE_       ,							strides=SCREAMING_SNAKE_CASE_       ,							padding='''VALID'''       ,							groups=SCREAMING_SNAKE_CASE_       ,							use_bias=SCREAMING_SNAKE_CASE_       ,							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 						_lowercase  (       self  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       )       ->     Any:
					lowercase_       =					self.convolution(self.padding(SCREAMING_SNAKE_CASE_       )       )
					lowercase_       =					self.normalization(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.activation(SCREAMING_SNAKE_CASE_       )
					return hidden_state

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



	"""simple docstring"""




	def __init__(       self  :			str       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							**SCREAMING_SNAKE_CASE_  :			str       )       ->     Any:
					super().__init__(**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					config.num_channels
					lowercase_       =					TFRegNetConvLayer(
					    out_channels=config.embedding_size       ,							kernel_size=3       ,							stride=2       ,							activation=config.hidden_act       ,							name='''embedder'''       ,							)




	def 						_lowercase  (       self  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			str       )       ->     Optional[int]:
					lowercase_       =					shape_list(SCREAMING_SNAKE_CASE_       )[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(SCREAMING_SNAKE_CASE_       ,							perm=(0, 2, 3, 1)       )
					lowercase_       =					self.embedder(SCREAMING_SNAKE_CASE_       )
					return hidden_state

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



	"""simple docstring"""




	def __init__(       self  :			int       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int = 2       ,							**SCREAMING_SNAKE_CASE_  :			List[str]       )       ->     Union[str, Any]:
					super().__init__(**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					tf.keras.layers.ConvaD(
					    filters=SCREAMING_SNAKE_CASE_       ,							kernel_size=1       ,							strides=SCREAMING_SNAKE_CASE_       ,							use_bias=SCREAMING_SNAKE_CASE_       ,							name='''convolution'''       )
					lowercase_       =					tf.keras.layers.BatchNormalization(epsilon=1e-5       ,							momentum=0.9       ,							name='''normalization'''       )




	def 						_lowercase  (       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor       ,							SCREAMING_SNAKE_CASE_  :			bool = False       )       ->     tf.Tensor:
					return self.normalization(self.convolution(SCREAMING_SNAKE_CASE_       )       ,							training=SCREAMING_SNAKE_CASE_       )

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



	"""simple docstring"""




	def __init__(       self  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int       ,							**SCREAMING_SNAKE_CASE_  :			int       )       ->     Optional[Any]:
					super().__init__(**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					tf.keras.layers.GlobalAveragePoolingaD(keepdims=SCREAMING_SNAKE_CASE_       ,							name='''pooler'''       )
					lowercase_       =					[
					    tf.keras.layers.ConvaD(filters=SCREAMING_SNAKE_CASE_       ,							kernel_size=1       ,							activation='''relu'''       ,							name='''attention.0'''       ),
					    tf.keras.layers.ConvaD(filters=SCREAMING_SNAKE_CASE_       ,							kernel_size=1       ,							activation='''sigmoid'''       ,							name='''attention.2'''       ),
					]




	def 						_lowercase  (       self  :			str       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       )       ->     Dict:
					# [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
					lowercase_       =					self.pooler(SCREAMING_SNAKE_CASE_       )
					for layer_module in self.attention:
									lowercase_       =					layer_module(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					hidden_state * pooled
					return hidden_state

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



	"""simple docstring"""




	def __init__(       self  :			str       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int = 1       ,							**SCREAMING_SNAKE_CASE_  :			Union[str, Any]       )       ->     List[str]:
					super().__init__(**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					in_channels != out_channels or stride != 1
					lowercase_       =					max(1       ,							out_channels // config.groups_width       )
					lowercase_       =					(
					    TFRegNetShortCut(SCREAMING_SNAKE_CASE_       ,							stride=SCREAMING_SNAKE_CASE_       ,							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(SCREAMING_SNAKE_CASE_       ,							kernel_size=1       ,							activation=config.hidden_act       ,							name='''layer.0'''       ),
					    TFRegNetConvLayer(
					        SCREAMING_SNAKE_CASE_       ,							stride=SCREAMING_SNAKE_CASE_       ,							groups=SCREAMING_SNAKE_CASE_       ,							activation=config.hidden_act       ,							name='''layer.1'''       ),
					    TFRegNetConvLayer(SCREAMING_SNAKE_CASE_       ,							kernel_size=1       ,							activation=SCREAMING_SNAKE_CASE_       ,							name='''layer.2'''       ),
					]
					lowercase_       =					ACTaFN[config.hidden_act]




	def 						_lowercase  (       self  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       )       ->     Any:
					lowercase_       =					hidden_state
					for layer_module in self.layers:
									lowercase_       =					layer_module(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.shortcut(SCREAMING_SNAKE_CASE_       )
					hidden_state += residual
					lowercase_       =					self.activation(SCREAMING_SNAKE_CASE_       )
					return hidden_state

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



	"""simple docstring"""




	def __init__(       self  :			Any       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int = 1       ,							**SCREAMING_SNAKE_CASE_  :			List[str]       )       ->     int:
					super().__init__(**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					in_channels != out_channels or stride != 1
					lowercase_       =					max(1       ,							out_channels // config.groups_width       )
					lowercase_       =					(
					    TFRegNetShortCut(SCREAMING_SNAKE_CASE_       ,							stride=SCREAMING_SNAKE_CASE_       ,							name='''shortcut'''       )
					    if should_apply_shortcut
					    else tf.keras.layers.Activation('''linear'''       ,							name='''shortcut'''       )
					)
					lowercase_       =					[
					    TFRegNetConvLayer(SCREAMING_SNAKE_CASE_       ,							kernel_size=1       ,							activation=config.hidden_act       ,							name='''layer.0'''       ),
					    TFRegNetConvLayer(
					        SCREAMING_SNAKE_CASE_       ,							stride=SCREAMING_SNAKE_CASE_       ,							groups=SCREAMING_SNAKE_CASE_       ,							activation=config.hidden_act       ,							name='''layer.1'''       ),
					    TFRegNetSELayer(SCREAMING_SNAKE_CASE_       ,							reduced_channels=int(round(in_channels / 4       )       )       ,							name='''layer.2'''       ),
					    TFRegNetConvLayer(SCREAMING_SNAKE_CASE_       ,							kernel_size=1       ,							activation=SCREAMING_SNAKE_CASE_       ,							name='''layer.3'''       ),
					]
					lowercase_       =					ACTaFN[config.hidden_act]




	def 						_lowercase  (       self  :			int       ,							SCREAMING_SNAKE_CASE_  :			Dict       )       ->     Optional[Any]:
					lowercase_       =					hidden_state
					for layer_module in self.layers:
									lowercase_       =					layer_module(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.shortcut(SCREAMING_SNAKE_CASE_       )
					hidden_state += residual
					lowercase_       =					self.activation(SCREAMING_SNAKE_CASE_       )
					return hidden_state

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



	"""simple docstring"""




	def __init__(       self  :			str       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int = 2       ,							SCREAMING_SNAKE_CASE_  :			int = 2       ,							**SCREAMING_SNAKE_CASE_  :			List[Any]       )       ->     List[str]:
					super().__init__(**SCREAMING_SNAKE_CASE_       )

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




	def 						_lowercase  (       self  :			Tuple       ,							SCREAMING_SNAKE_CASE_  :			int       )       ->     int:
					for layer_module in self.layers:
									lowercase_       =					layer_module(SCREAMING_SNAKE_CASE_       )
					return hidden_state

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



	"""simple docstring"""




	def __init__(       self  :			Union[str, Any]       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							**SCREAMING_SNAKE_CASE_  :			Dict       )       ->     Tuple:
					super().__init__(**SCREAMING_SNAKE_CASE_       )
					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(
					        SCREAMING_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'''       ,							)       )
					lowercase_       =					zip(config.hidden_sizes       ,							config.hidden_sizes[1:]       )
					for i, ((in_channels, out_channels), depth) in enumerate(zip(SCREAMING_SNAKE_CASE_       ,							config.depths[1:]       )       ):
									self.stages.append(TFRegNetStage(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							depth=SCREAMING_SNAKE_CASE_       ,							name=f'''stages.{i+1}'''       )       )




	def 						_lowercase  (       self  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor       ,							SCREAMING_SNAKE_CASE_  :			bool = False       ,							SCREAMING_SNAKE_CASE_  :			bool = True       )       ->     TFBaseModelOutputWithNoAttention:
					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(SCREAMING_SNAKE_CASE_       )

					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=SCREAMING_SNAKE_CASE_       ,							hidden_states=SCREAMING_SNAKE_CASE_       )

@keras_serializable
class  lowercase__(				tf.keras.layers.Layer							):



	"""simple docstring"""




	a			:str         =	RegNetConfig
	def __init__(       self  :			str       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							**SCREAMING_SNAKE_CASE_  :			List[str]       )       ->     Any:
					super().__init__(**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					config
					lowercase_       =					TFRegNetEmbeddings(SCREAMING_SNAKE_CASE_       ,							name='''embedder'''       )
					lowercase_       =					TFRegNetEncoder(SCREAMING_SNAKE_CASE_       ,							name='''encoder'''       )
					lowercase_       =					tf.keras.layers.GlobalAveragePoolingaD(keepdims=SCREAMING_SNAKE_CASE_       ,							name='''pooler'''       )




	@unpack_inputs
	def 						_lowercase  (       self  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							SCREAMING_SNAKE_CASE_  :			bool = False       ,							)       ->     TFBaseModelOutputWithPoolingAndNoAttention:
					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(SCREAMING_SNAKE_CASE_       ,							training=SCREAMING_SNAKE_CASE_       )

					lowercase_       =					self.encoder(
					    SCREAMING_SNAKE_CASE_       ,							output_hidden_states=SCREAMING_SNAKE_CASE_       ,							return_dict=SCREAMING_SNAKE_CASE_       ,							training=SCREAMING_SNAKE_CASE_       )

					lowercase_       =					encoder_outputs[0]
					lowercase_       =					self.pooler(SCREAMING_SNAKE_CASE_       )

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

					# Change the other hidden state outputs to NCHW as well
					if output_hidden_states:
									lowercase_       =					tuple([tf.transpose(SCREAMING_SNAKE_CASE_       ,							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=SCREAMING_SNAKE_CASE_       ,							pooler_output=SCREAMING_SNAKE_CASE_       ,							hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states       ,							)

class  lowercase__(				UpperCAmelCase							):



	"""simple docstring"""




	a			:Tuple         =	RegNetConfig
	a			:Any         =	'regnet'
	a			:List[str]         =	'pixel_values'
	@property
	def 						_lowercase  (       self  :			List[str]       )       ->     str:
					return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_2_4, 2_2_4)       ,							dtype=tf.floataa       )}

__a  			=					r'\n    Parameters:\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        config ([`RegNetConfig`]): 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            [`ConveNextImageProcessor.__call__`] for details.\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'
@add_start_docstrings(
    'The bare RegNet model outputting raw features without any specific head on top.' ,						UpperCAmelCase ,						)
class  lowercase__(				UpperCAmelCase							):



	"""simple docstring"""




	def __init__(       self  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							*SCREAMING_SNAKE_CASE_  :			Any       ,							**SCREAMING_SNAKE_CASE_  :			str       )       ->     List[str]:
					super().__init__(SCREAMING_SNAKE_CASE_       ,							*SCREAMING_SNAKE_CASE_       ,							**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					TFRegNetMainLayer(SCREAMING_SNAKE_CASE_       ,							name='''regnet'''       )




	@unpack_inputs
	@add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_       )
	@add_code_sample_docstrings(
	    checkpoint=_CHECKPOINT_FOR_DOC       ,							output_type=SCREAMING_SNAKE_CASE_       ,							config_class=_CONFIG_FOR_DOC       ,							modality='''vision'''       ,							expected_output=_EXPECTED_OUTPUT_SHAPE       ,							)
	def 						_lowercase  (       self  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]=False       ,							)       ->     Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:
					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=SCREAMING_SNAKE_CASE_       ,							output_hidden_states=SCREAMING_SNAKE_CASE_       ,							return_dict=SCREAMING_SNAKE_CASE_       ,							training=SCREAMING_SNAKE_CASE_       ,							)
					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    ' ,						UpperCAmelCase ,						)
class  lowercase__(				UpperCAmelCase ,						UpperCAmelCase							):



	"""simple docstring"""




	def __init__(       self  :			Union[str, Any]       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							*SCREAMING_SNAKE_CASE_  :			str       ,							**SCREAMING_SNAKE_CASE_  :			int       )       ->     Union[str, Any]:
					super().__init__(SCREAMING_SNAKE_CASE_       ,							*SCREAMING_SNAKE_CASE_       ,							**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					config.num_labels
					lowercase_       =					TFRegNetMainLayer(SCREAMING_SNAKE_CASE_       ,							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(SCREAMING_SNAKE_CASE_       )
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT       ,							output_type=SCREAMING_SNAKE_CASE_       ,							config_class=_CONFIG_FOR_DOC       ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT       ,							)
	def 						_lowercase  (       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor = None       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor = None       ,							SCREAMING_SNAKE_CASE_  :			bool = None       ,							SCREAMING_SNAKE_CASE_  :			bool = None       ,							SCREAMING_SNAKE_CASE_  :			Dict=False       ,							)       ->     Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
					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(
					    SCREAMING_SNAKE_CASE_       ,							output_hidden_states=SCREAMING_SNAKE_CASE_       ,							return_dict=SCREAMING_SNAKE_CASE_       ,							training=SCREAMING_SNAKE_CASE_       )

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

					lowercase_       =					self.classifier[0](SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.classifier[1](SCREAMING_SNAKE_CASE_       )

					lowercase_       =					None if labels is None else self.hf_compute_loss(labels=SCREAMING_SNAKE_CASE_       ,							logits=SCREAMING_SNAKE_CASE_       )

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

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




