
from __future__ import annotations

import unittest

import numpy as np

from transformers import OPTConfig, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import GPTaTokenizer, TFOPTForCausalLM, TFOPTModel


def 						UpperCAmelCase			(	a_		,   a_		,   a_=None		,   a_=None						)		->	Any:







				"""simple docstring"""



				if attention_mask is None:
								__A   			= tf.cast(tf.math.not_equal(a_		,   config.pad_token_id						)		,   tf.inta						)
				return {"input_ids": input_ids, "attention_mask": attention_mask}





@require_tf
class 							UpperCAmelCase      :




	'''simple docstring'''





	snake_case_					    =     OPTConfig
	snake_case_					    =     {}
	snake_case_					    =     "gelu"






	def __init__(       self					: Tuple							,A					: Tuple							,A					: Optional[Any]=13							,A					: Union[str, Any]=7							,A					: Tuple=True							,A					: Optional[int]=False							,A					: List[str]=99							,A					: str=16							,A					: Optional[Any]=2							,A					: List[str]=4							,A					: Optional[int]=4							,A					: str="gelu"							,A					: Any=0.1							,A					: int=0.1							,A					: int=20							,A					: Tuple=2							,A					: Optional[int]=1							,A					: Union[str, Any]=0							,A					: str=16							,A					: Dict=16							,):
					__A   			= parent
					__A   			= batch_size
					__A   			= seq_length
					__A   			= is_training
					__A   			= use_labels
					__A   			= vocab_size
					__A   			= hidden_size
					__A   			= num_hidden_layers
					__A   			= num_attention_heads
					__A   			= intermediate_size
					__A   			= hidden_act
					__A   			= hidden_dropout_prob
					__A   			= attention_probs_dropout_prob
					__A   			= max_position_embeddings
					__A   			= eos_token_id
					__A   			= pad_token_id
					__A   			= bos_token_id
					__A   			= embed_dim
					__A   			= word_embed_proj_dim
					__A   			= False






	def 				UpperCamelCase_  (       self					: List[str] ):
					__A   			= ids_tensor([self.batch_size, self.seq_length - 1]							,self.vocab_size )
					__A   			= tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size )							,1 )
					__A   			= tf.concat([input_ids, eos_tensor]							,axis=1 )

					__A   			= self.config_cls(
					    vocab_size=self.vocab_size							,hidden_size=self.hidden_size							,num_hidden_layers=self.num_hidden_layers							,num_attention_heads=self.num_attention_heads							,ffn_dim=self.intermediate_size							,dropout=self.hidden_dropout_prob							,attention_dropout=self.attention_probs_dropout_prob							,max_position_embeddings=self.max_position_embeddings							,eos_token_id=self.eos_token_id							,bos_token_id=self.bos_token_id							,pad_token_id=self.pad_token_id							,embed_dim=self.embed_dim							,word_embed_proj_dim=self.word_embed_proj_dim							,is_encoder_decoder=A							,**self.config_updates							,)
					__A   			= prepare_opt_inputs_dict(A							,A )
					return config, inputs_dict







	def 				UpperCamelCase_  (       self					: Dict							,A					: int							,A					: Optional[int] ):
					__A   			= TFOPTModel(config=A )
					__A   			= inputs_dict["input_ids"]

					__A   			= input_ids[:1, :]
					__A   			= inputs_dict["attention_mask"][:1, :]
					__A   			= 1

					# first forward pass
					__A   			= model(A							,attention_mask=A							,use_cache=A )

					__A			,							__A   			= outputs.to_tuple()

					# create hypothetical next token and extent to next_input_ids
					__A   			= ids_tensor((self.batch_size, 3)							,config.vocab_size )
					__A   			= tf.cast(ids_tensor((self.batch_size, 3)							,2 )							,tf.inta )

					# append to next input_ids and
					__A   			= tf.concat([input_ids, next_tokens]							,axis=-1 )
					__A   			= tf.concat([attention_mask, next_attn_mask]							,axis=-1 )

					__A   			= model(A							,attention_mask=A )[0]
					__A   			= model(A							,attention_mask=A							,past_key_values=A )[0]

					self.parent.assertEqual(next_tokens.shape[1]							,output_from_past.shape[1] )

					# select random slice
					__A   			= int(ids_tensor((1,)							,output_from_past.shape[-1] ) )
					__A   			= output_from_no_past[:, -3:, random_slice_idx]
					__A   			= output_from_past[:, :, random_slice_idx]

					# test that outputs are equal for slice
					tf.debugging.assert_near(A							,A							,rtol=1E-3 )





@require_tf
class 							UpperCAmelCase      (  __SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			unittest.TestCase  ):




	'''simple docstring'''





	snake_case_					    =     (TFOPTModel, TFOPTForCausalLM) if is_tf_available() else ()
	snake_case_					    =     (TFOPTForCausalLM,) if is_tf_available() else ()
	snake_case_					    =     (
	    {"feature-extraction": TFOPTModel, "text-generation": TFOPTForCausalLM} if is_tf_available() else {}
	)
	snake_case_					    =     False
	snake_case_					    =     False
	snake_case_					    =     False
	snake_case_					    =     10






	def 				UpperCamelCase_  (       self					: Union[str, Any] ):
					__A   			= TFOPTModelTester(self )
					__A   			= ConfigTester(self							,config_class=A )






	def 				UpperCamelCase_  (       self					: Dict ):
					self.config_tester.run_common_tests()






	def 				UpperCamelCase_  (       self					: Optional[int] ):
					__A   			= self.model_tester.prepare_config_and_inputs_for_common()
					self.model_tester.check_decoder_model_past_large_inputs(*A )







	def 				UpperCamelCase_  (       self					: Any ):
					__A			,							__A   			= self.model_tester.prepare_config_and_inputs_for_common()

					def _get_word_embedding_weight(A					: Optional[int]							,A					: Optional[Any] ):
									if hasattr(A							,"weight" ):
													return embedding_layer.weight
									else:
													# Here we build the word embeddings weights if not exists.
													# And then we retry to get the attribute once built.
													model.build()
													if hasattr(A							,"weight" ):
																	return embedding_layer.weight
													else:
																	return None

					for model_class in self.all_model_classes:
									for size in [config.vocab_size - 10, config.vocab_size + 10]:
													# build the embeddings
													__A   			= model_class(config=A )
													__A   			= _get_word_embedding_weight(A							,model.get_input_embeddings() )
													__A   			= _get_word_embedding_weight(A							,model.get_output_embeddings() )

													# reshape the embeddings
													model.resize_token_embeddings(A )
													__A   			= _get_word_embedding_weight(A							,model.get_input_embeddings() )
													__A   			= _get_word_embedding_weight(A							,model.get_output_embeddings() )

													# check that the resized embeddings size matches the desired size.
													__A   			= size if size is not None else config.vocab_size

													self.assertEqual(new_input_embeddings.shape[0]							,A )

													# check that weights remain the same after resizing
													__A   			= True
													for pa, pa in zip(old_input_embeddings.value()							,new_input_embeddings.value() ):
																	if tf.math.reduce_sum(tf.math.abs(pa - pa ) ) > 0:
																					__A   			= False
													self.assertTrue(A )

													if old_output_embeddings is not None and new_output_embeddings is not None:
																	self.assertEqual(new_output_embeddings.shape[0]							,A )

																	__A   			= True
																	for pa, pa in zip(old_output_embeddings.value()							,new_output_embeddings.value() ):
																					if tf.math.reduce_sum(tf.math.abs(pa - pa ) ) > 0:
																									__A   			= False
																	self.assertTrue(A )





def 						UpperCAmelCase			(	a_						)		->	List[str]:







				"""simple docstring"""



				return tf.constant(a_		,   dtype=tf.intaa						)





@require_tf
class 							UpperCAmelCase      (  unittest.TestCase  ):




	'''simple docstring'''





	snake_case_					    =     99






	def 				UpperCamelCase_  (       self					: Tuple ):
					__A   			= tf.ones((4, 1)							,dtype=tf.intaa ) * 2
					__A   			= tf.concat([ids_tensor((4, 6)							,self.vocab_size - 3 ) + 3, eos_column_vector]							,axis=1 )
					__A   			= input_ids.shape[0]
					__A   			= OPTConfig(
					    vocab_size=self.vocab_size							,hidden_size=24							,num_hidden_layers=2							,num_attention_heads=2							,ffn_dim=32							,max_position_embeddings=48							,eos_token_id=2							,pad_token_id=1							,bos_token_id=0							,)
					return config, input_ids, batch_size





@require_sentencepiece
@require_tf
class 							UpperCAmelCase      (  unittest.TestCase  ):




	'''simple docstring'''





	@slow
	def 				UpperCamelCase_  (       self					: List[str] ):
					__A   			= TFOPTModel.from_pretrained("facebook/opt-350m" )
					__A   			= _long_tensor([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]] )
					__A   			= tf.not_equal(A							,model.config.pad_token_id )
					with tf.GradientTape():
									__A   			= model(input_ids=A							,attention_mask=A ).last_hidden_state
					__A   			= (1, 11, 5_12)
					self.assertEqual(output.shape							,A )
					__A   			= tf.constant(
					    [[-0.28_73, -1.92_18, -0.30_33], [-1.27_10, -0.13_38, -0.19_02], [0.40_95, 0.12_14, -1.31_21]] )
					self.assertTrue(np.allclose(output[:, :3, :3]							,A							,atol=4E-3 ) )

					__A   			= tf.function(A							,jit_compile=A )
					__A   			= xla_generate(A							,A )[0]
					self.assertTrue(np.allclose(output[:, :3, :3]							,A							,atol=4E-2 ) )





@require_tf
@slow
class 							UpperCAmelCase      (  unittest.TestCase  ):




	'''simple docstring'''





	def 				UpperCamelCase_  (       self					: Optional[Any] ):
					super().setUp()
					__A   			= "facebook/opt-350m"







	def 				UpperCamelCase_  (       self					: Optional[int] ):
					__A   			= TFOPTForCausalLM.from_pretrained(self.path_model )
					__A   			= GPTaTokenizer.from_pretrained(self.path_model )

					__A   			= [
					    "Today is a beautiful day and I want to",
					    "In the city of",
					    "Paris is the capital of France and",
					    "Computers and mobile phones have taken",
					]
					# verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
					__A   			= tokenizer(A							,return_tensors="tf"							,padding=A							,add_special_tokens=A )
					__A   			= tf.math.reduce_mean(model(inputs.input_ids							,attention_mask=inputs.attention_mask )[0]							,axis=-1 )
					__A   			= tf.constant(
					    [
					        [1.38_51, -13.89_23, -10.52_29, -10.75_33, -0.23_09, -10.23_84, -0.53_65, -9.09_47, -5.16_70],
					        [-4.70_73, -10.62_76, -3.94_15, -21.52_42, -0.28_22, -0.28_22, -0.28_22, -0.28_22, -0.28_22],
					        [0.62_47, -3.42_29, -8.91_79, -1.42_97, -14.16_50, 1.41_46, -9.02_18, -0.27_03, -0.27_03],
					        [6.47_83, -1.99_13, -10.79_26, -2.33_36, 1.50_92, -0.99_74, -6.82_13, 1.34_77, 1.34_77],
					    ] )
					self.assertTrue(np.allclose(A							,A							,atol=1E-4 ) )

					__A   			= tf.function(A							,jit_compile=A )
					__A   			= tf.math.reduce_mean(xla_generate(inputs.input_ids							,attention_mask=inputs.attention_mask )[0]							,axis=-1 )
					self.assertTrue(np.allclose(A							,A							,atol=1E-4 ) )





@require_tf
@slow
class 							UpperCAmelCase      (  unittest.TestCase  ):




	'''simple docstring'''





	@property
	def 				UpperCamelCase_  (       self					: int ):
					return [
					    "Today is a beautiful day and I want",
					    "In the city of",
					    "Paris is the capital of France and",
					    "Computers and mobile phones have taken",
					]






	def 				UpperCamelCase_  (       self					: Optional[int] ):
					__A   			= "facebook/opt-125m"

					__A   			= [
					    "Today is a beautiful day and I want to",
					    "In the city of New York, the city",
					    "Paris is the capital of France and the capital",
					    "Computers and mobile phones have taken over the",
					]

					__A   			= []
					__A   			= GPTaTokenizer.from_pretrained(A )
					__A   			= TFOPTForCausalLM.from_pretrained(A )

					for prompt in self.prompts:
									__A   			= tokenizer(A							,return_tensors="tf" ).input_ids

									__A   			= model.generate(A							,max_length=10 )

									__A   			= tokenizer.batch_decode(A							,skip_special_tokens=A )
									predicted_outputs += generated_string

					self.assertListEqual(A							,A )






	def 				UpperCamelCase_  (       self					: int ):
					__A   			= "facebook/opt-350m"

					__A   			= GPTaTokenizer.from_pretrained(A )
					__A   			= TFOPTForCausalLM.from_pretrained(A )

					__A   			= "left"

					# use different length sentences to test batching
					__A   			= [
					    "Hello, my dog is a little",
					    "Today, I",
					]

					__A   			= tokenizer(A							,return_tensors="tf"							,padding=A )
					__A   			= inputs["input_ids"]

					__A   			= model.generate(input_ids=A							,attention_mask=inputs["attention_mask"] )

					__A   			= tokenizer(sentences[0]							,return_tensors="tf" ).input_ids
					__A   			= model.generate(input_ids=A )

					__A   			= inputs_non_padded.shape[-1] - tf.math.reduce_sum(
					    tf.cast(inputs["attention_mask"][-1]							,tf.intaa ) )
					__A   			= tokenizer(sentences[1]							,return_tensors="tf" ).input_ids
					__A   			= model.generate(input_ids=A							,max_length=model.config.max_length - num_paddings )

					__A   			= tokenizer.batch_decode(A							,skip_special_tokens=A )
					__A   			= tokenizer.decode(output_non_padded[0]							,skip_special_tokens=A )
					__A   			= tokenizer.decode(output_padded[0]							,skip_special_tokens=A )

					__A   			= [
					    "Hello, my dog is a little bit of a dork.\nI'm a little bit",
					    "Today, I was in the middle of a conversation with a friend about the",
					]
					self.assertListEqual(A							,A )
					self.assertListEqual(A							,[non_padded_sentence, padded_sentence] )







	def 				UpperCamelCase_  (       self					: Dict ):
					__A   			= "facebook/opt-350m"

					__A   			= [
					    "Today is a beautiful day and I want to",
					    "In the city of San Francisco, the city",
					    "Paris is the capital of France and the capital",
					    "Computers and mobile phones have taken over the",
					]

					__A   			= []
					__A   			= GPTaTokenizer.from_pretrained(A )
					__A   			= TFOPTForCausalLM.from_pretrained(A )

					for prompt in self.prompts:
									__A   			= tokenizer(A							,return_tensors="tf" ).input_ids

									__A   			= model.generate(A							,max_length=10 )

									__A   			= tokenizer.batch_decode(A							,skip_special_tokens=A )
									predicted_outputs += generated_string

					self.assertListEqual(A							,A )


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 )
