
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


SCREAMING_SNAKE_CASE							:str							 =				logging.get_logger(__name__)

# General docstring
SCREAMING_SNAKE_CASE							:str							 =				'RegNetConfig'

# Base docstring
SCREAMING_SNAKE_CASE							:List[str]							 =				'facebook/regnet-y-040'
SCREAMING_SNAKE_CASE							:Union[str, Any]							 =				[1, 1088, 7, 7]

# Image classification docstring
SCREAMING_SNAKE_CASE							:Optional[int]							 =				'facebook/regnet-y-040'
SCREAMING_SNAKE_CASE							:Any							 =				'tabby, tabby cat'

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





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




	'''simple docstring'''





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







	def 				UpperCamelCase_  (       self					: List[Any]							,A					: Any ):
					__A   			= self.convolution(self.padding(A ) )
					__A   			= self.normalization(A )
					__A   			= self.activation(A )
					return hidden_state





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




	'''simple docstring'''





	def __init__(       self					: Tuple							,A					: RegNetConfig							,**A					: str ):
					super().__init__(**A )
					__A   			= config.num_channels
					__A   			= TFRegNetConvLayer(
					    out_channels=config.embedding_size							,kernel_size=3							,stride=2							,activation=config.hidden_act							,name="embedder"							,)







	def 				UpperCamelCase_  (       self					: Tuple							,A					: Optional[Any] ):
					__A   			= shape_list(A )[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)
					__A   			= tf.transpose(A							,perm=(0, 2, 3, 1) )
					__A   			= self.embedder(A )
					return hidden_state





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




	'''simple docstring'''





	def __init__(       self					: Optional[int]							,A					: int							,A					: int = 2							,**A					: Tuple ):
					super().__init__(**A )
					__A   			= tf.keras.layers.ConvaD(
					    filters=A							,kernel_size=1							,strides=A							,use_bias=A							,name="convolution" )
					__A   			= tf.keras.layers.BatchNormalization(epsilon=1E-5							,momentum=0.9							,name="normalization" )







	def 				UpperCamelCase_  (       self					: Union[str, Any]							,A					: tf.Tensor							,A					: bool = False ):
					return self.normalization(self.convolution(A )							,training=A )





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




	'''simple docstring'''





	def __init__(       self					: Dict							,A					: int							,A					: int							,**A					: str ):
					super().__init__(**A )
					__A   			= tf.keras.layers.GlobalAveragePoolingaD(keepdims=A							,name="pooler" )
					__A   			= [
					    tf.keras.layers.ConvaD(filters=A							,kernel_size=1							,activation="relu"							,name="attention.0" ),
					    tf.keras.layers.ConvaD(filters=A							,kernel_size=1							,activation="sigmoid"							,name="attention.2" ),
					]







	def 				UpperCamelCase_  (       self					: Dict							,A					: List[Any] ):
					# [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
					__A   			= self.pooler(A )
					for layer_module in self.attention:
									__A   			= layer_module(A )
					__A   			= hidden_state * pooled
					return hidden_state





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




	'''simple docstring'''





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







	def 				UpperCamelCase_  (       self					: int							,A					: Optional[int] ):
					__A   			= hidden_state
					for layer_module in self.layers:
									__A   			= layer_module(A )
					__A   			= self.shortcut(A )
					hidden_state += residual
					__A   			= self.activation(A )
					return hidden_state





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




	'''simple docstring'''





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







	def 				UpperCamelCase_  (       self					: Dict							,A					: Any ):
					__A   			= hidden_state
					for layer_module in self.layers:
									__A   			= layer_module(A )
					__A   			= self.shortcut(A )
					hidden_state += residual
					__A   			= self.activation(A )
					return hidden_state





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




	'''simple docstring'''





	def __init__(       self					: List[str]							,A					: RegNetConfig							,A					: int							,A					: int							,A					: int = 2							,A					: int = 2							,**A					: Optional[int] ):
					super().__init__(**A )

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







	def 				UpperCamelCase_  (       self					: Any							,A					: List[str] ):
					for layer_module in self.layers:
									__A   			= layer_module(A )
					return hidden_state





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




	'''simple docstring'''





	def __init__(       self					: Any							,A					: RegNetConfig							,**A					: List[str] ):
					super().__init__(**A )
					__A   			= []
					# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
					self.stages.append(
					    TFRegNetStage(
					        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"							,) )
					__A   			= zip(config.hidden_sizes							,config.hidden_sizes[1:] )
					for i, ((in_channels, out_channels), depth) in enumerate(zip(A							,config.depths[1:] ) ):
									self.stages.append(TFRegNetStage(A							,A							,A							,depth=A							,name=f'''stages.{i+1}''' ) )







	def 				UpperCamelCase_  (       self					: List[str]							,A					: tf.Tensor							,A					: bool = False							,A					: bool = True ):
					__A   			= () if output_hidden_states else None

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

									__A   			= stage_module(A )

					if output_hidden_states:
									__A   			= 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 )





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




	'''simple docstring'''





	snake_case_					    =     RegNetConfig






	def __init__(       self					: int							,A					: Optional[int]							,**A					: Dict ):
					super().__init__(**A )
					__A   			= config
					__A   			= TFRegNetEmbeddings(A							,name="embedder" )
					__A   			= TFRegNetEncoder(A							,name="encoder" )
					__A   			= tf.keras.layers.GlobalAveragePoolingaD(keepdims=A							,name="pooler" )







	@unpack_inputs
	def 				UpperCamelCase_  (       self					: Tuple							,A					: tf.Tensor							,A					: Optional[bool] = None							,A					: Optional[bool] = None							,A					: bool = False							,):
					__A   			= (
					    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
					)
					__A   			= return_dict if return_dict is not None else self.config.use_return_dict

					__A   			= self.embedder(A							,training=A )

					__A   			= self.encoder(
					    A							,output_hidden_states=A							,return_dict=A							,training=A )

					__A   			= encoder_outputs[0]
					__A   			= self.pooler(A )

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

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





class 							UpperCAmelCase      (  __SCREAMING_SNAKE_CASE  ):




	'''simple docstring'''





	snake_case_					    =     RegNetConfig
	snake_case_					    =     "regnet"
	snake_case_					    =     "pixel_values"






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





SCREAMING_SNAKE_CASE							:Dict							 =				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'

SCREAMING_SNAKE_CASE							:Dict							 =				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."			,			__SCREAMING_SNAKE_CASE			,			)
class 							UpperCAmelCase      (  __SCREAMING_SNAKE_CASE  ):




	'''simple docstring'''





	def __init__(       self					: List[Any]							,A					: RegNetConfig							,*A					: List[Any]							,**A					: str ):
					super().__init__(A							,*A							,**A )
					__A   			= TFRegNetMainLayer(A							,name="regnet" )







	@unpack_inputs
	@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							,)
	def 				UpperCamelCase_  (       self					: Tuple							,A					: tf.Tensor							,A					: Optional[bool] = None							,A					: Optional[bool] = None							,A					: int=False							,):
					__A   			= (
					    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
					)
					__A   			= return_dict if return_dict is not None else self.config.use_return_dict

					__A   			= self.regnet(
					    pixel_values=A							,output_hidden_states=A							,return_dict=A							,training=A							,)
					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    "			,			__SCREAMING_SNAKE_CASE			,			)
class 							UpperCAmelCase      (  __SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE  ):




	'''simple docstring'''





	def __init__(       self					: Optional[int]							,A					: RegNetConfig							,*A					: str							,**A					: Tuple ):
					super().__init__(A							,*A							,**A )
					__A   			= config.num_labels
					__A   			= TFRegNetMainLayer(A							,name="regnet" )
					# classification head
					__A   			= [
					    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(A )
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT							,output_type=A							,config_class=_CONFIG_FOR_DOC							,expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							,)
	def 				UpperCamelCase_  (       self					: List[str]							,A					: tf.Tensor = None							,A					: tf.Tensor = None							,A					: bool = None							,A					: bool = None							,A					: Union[str, Any]=False							,):
					__A   			= (
					    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
					)
					__A   			= return_dict if return_dict is not None else self.config.use_return_dict

					__A   			= self.regnet(
					    A							,output_hidden_states=A							,return_dict=A							,training=A )

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

					__A   			= self.classifier[0](A )
					__A   			= self.classifier[1](A )

					__A   			= None if labels is None else self.hf_compute_loss(labels=A							,logits=A )

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

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







from ...configuration_utils import PretrainedConfig
from ...utils import logging


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

_snake_case  :    Any						    =  {
    'google/pegasus-large': 'https://huggingface.co/google/pegasus-large/resolve/main/config.json',
    # See all PEGASUS models at https://huggingface.co/models?filter=pegasus
}



class 				_UpperCAmelCase    (    _UpperCamelCase    ):

  """simple docstring"""


  a_        = """pegasus"""
  a_        = ["""past_key_values"""]
  a_        = {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






  def __init__(				self :		List[str]					,  lowerCAmelCase_ :		Union[str, Any]=5_0_2_6_5					,  lowerCAmelCase_ :		Union[str, Any]=1_0_2_4					,  lowerCAmelCase_ :		Union[str, Any]=1_2					,  lowerCAmelCase_ :		Dict=4_0_9_6					,  lowerCAmelCase_ :		str=1_6					,  lowerCAmelCase_ :		List[Any]=1_2					,  lowerCAmelCase_ :		Union[str, Any]=4_0_9_6					,  lowerCAmelCase_ :		Union[str, Any]=1_6					,  lowerCAmelCase_ :		Any=0.0					,  lowerCAmelCase_ :		Optional[int]=0.0					,  lowerCAmelCase_ :		List[Any]=True					,  lowerCAmelCase_ :		Any=True					,  lowerCAmelCase_ :		Dict="gelu"					,  lowerCAmelCase_ :		Optional[int]=1_0_2_4					,  lowerCAmelCase_ :		List[str]=0.1					,  lowerCAmelCase_ :		str=0.0					,  lowerCAmelCase_ :		Union[str, Any]=0.0					,  lowerCAmelCase_ :		Optional[int]=0.02					,  lowerCAmelCase_ :		Tuple=0					,  lowerCAmelCase_ :		Tuple=False					,  lowerCAmelCase_ :		Union[str, Any]=0					,  lowerCAmelCase_ :		int=1					,  lowerCAmelCase_ :		Tuple=1					,  **lowerCAmelCase_ :		Tuple					,  )							->						List[str]:
     __lowerCAmelCase								=	vocab_size
     __lowerCAmelCase								=	max_position_embeddings
     __lowerCAmelCase								=	d_model
     __lowerCAmelCase								=	encoder_ffn_dim
     __lowerCAmelCase								=	encoder_layers
     __lowerCAmelCase								=	encoder_attention_heads
     __lowerCAmelCase								=	decoder_ffn_dim
     __lowerCAmelCase								=	decoder_layers
     __lowerCAmelCase								=	decoder_attention_heads
     __lowerCAmelCase								=	dropout
     __lowerCAmelCase								=	attention_dropout
     __lowerCAmelCase								=	activation_dropout
     __lowerCAmelCase								=	activation_function
     __lowerCAmelCase								=	init_std
     __lowerCAmelCase								=	encoder_layerdrop
     __lowerCAmelCase								=	decoder_layerdrop
     __lowerCAmelCase								=	use_cache
     __lowerCAmelCase								=	encoder_layers
     __lowerCAmelCase								=	scale_embedding  # scale factor will be sqrt(d_model) if True
     super().__init__(
         pad_token_id=lowerCAmelCase_					,  eos_token_id=lowerCAmelCase_					,  is_encoder_decoder=lowerCAmelCase_					,  decoder_start_token_id=lowerCAmelCase_					,  forced_eos_token_id=lowerCAmelCase_					,  **lowerCAmelCase_					,  )






  @property
  def 				lowercase		(				self :		List[Any]     )							->						int:
     return self.encoder_attention_heads






  @property
  def 				lowercase		(				self :		Optional[Any]     )							->						int:
     return self.d_model

