






"""simple docstring"""






import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
    import os

    # The slow tests are often failing with OOM error on GPU
    # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
    # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
    A_			  =    "platform"

    import jax
    import jax.numpy as jnp

    from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer




def 						_UpperCamelCase			(  A    ,       A    ,       A=None    ,       A=None			):







   if attention_mask is None:
      UpperCamelCase_     =np.where(input_ids != config.pad_token_id    ,       1    ,       0			)
   return {
       "input_ids": input_ids,
       "attention_mask": attention_mask,
   }
@require_flax
class 						__lowerCAmelCase   :







  '''simple docstring'''


  def __init__(      self:		Any				,			UpperCamelCase_:		List[Any]				,			UpperCamelCase_:		Tuple=13				,			UpperCamelCase_:		Optional[int]=7				,			UpperCamelCase_:		str=True				,			UpperCamelCase_:		str=False				,			UpperCamelCase_:		List[Any]=99				,			UpperCamelCase_:		Optional[Any]=16				,			UpperCamelCase_:		Union[str, Any]=2				,			UpperCamelCase_:		List[Any]=4				,			UpperCamelCase_:		Tuple=4				,			UpperCamelCase_:		Any="gelu"				,			UpperCamelCase_:		Any=0.1				,			UpperCamelCase_:		Tuple=0.1				,			UpperCamelCase_:		List[str]=20				,			UpperCamelCase_:		int=2				,			UpperCamelCase_:		Optional[int]=1				,			UpperCamelCase_:		Dict=0				,			UpperCamelCase_:		Tuple=16				,			UpperCamelCase_:		int=16				,			UpperCamelCase_:		int=0.02				,			):
     UpperCamelCase_     =parent
     UpperCamelCase_     =batch_size
     UpperCamelCase_     =seq_length
     UpperCamelCase_     =is_training
     UpperCamelCase_     =use_labels
     UpperCamelCase_     =vocab_size
     UpperCamelCase_     =hidden_size
     UpperCamelCase_     =num_hidden_layers
     UpperCamelCase_     =num_attention_heads
     UpperCamelCase_     =intermediate_size
     UpperCamelCase_     =hidden_act
     UpperCamelCase_     =hidden_dropout_prob
     UpperCamelCase_     =attention_probs_dropout_prob
     UpperCamelCase_     =max_position_embeddings
     UpperCamelCase_     =eos_token_id
     UpperCamelCase_     =pad_token_id
     UpperCamelCase_     =bos_token_id
     UpperCamelCase_     =embed_dim
     UpperCamelCase_     =word_embed_proj_dim
     UpperCamelCase_     =initializer_range
     UpperCamelCase_     =False
  def 					UpperCamelCase__   (      self:		Tuple					):
     UpperCamelCase_     =np.clip(ids_tensor([self.batch_size, self.seq_length - 1]				,			self.vocab_size					)				,			3				,			self.vocab_size					)
     UpperCamelCase_     =np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)				,			dtype=np.intaa					))				,			-1					)

     UpperCamelCase_     =OPTConfig(
         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				,			is_encoder_decoder=a_				,			word_embed_proj_dim=self.word_embed_proj_dim				,			initializer_range=self.initializer_range				,			use_cache=a_				,			)
     UpperCamelCase_     =prepare_opt_inputs_dict(a_				,			a_					)
     return config, inputs_dict
  def 					UpperCamelCase__   (      self:		Any					):
     UpperCamelCase_    ,							UpperCamelCase_     =self.prepare_config_and_inputs()
     return config, inputs_dict
  def 					UpperCamelCase__   (      self:		Any				,			UpperCamelCase_:		Optional[Any]				,			UpperCamelCase_:		List[str]				,			UpperCamelCase_:		Optional[Any]					):
     UpperCamelCase_     =20
     UpperCamelCase_     =model_class_name(a_					)

     UpperCamelCase_     =inputs_dict["input_ids"]
     UpperCamelCase_     =inputs_dict["attention_mask"]

     UpperCamelCase_     =model.init_cache(input_ids.shape[0]				,			a_					)
     UpperCamelCase_     =jnp.ones((input_ids.shape[0], max_length)				,			dtype="i4"					)

     UpperCamelCase_     =jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1					)[None, :]				,			(input_ids.shape[0], input_ids.shape[-1] - 1)				,			)
     UpperCamelCase_     =model(
         input_ids[:, :-1]				,			attention_mask=a_				,			past_key_values=a_				,			position_ids=a_				,			)

     UpperCamelCase_     =jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,			dtype="i4"					)
     UpperCamelCase_     =model(
         input_ids[:, -1:]				,			attention_mask=a_				,			past_key_values=outputs_cache.past_key_values				,			position_ids=a_				,			)

     UpperCamelCase_     =model(a_					)

     UpperCamelCase_     =np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])					)					)
     self.parent.assertTrue(diff < 1e-3				,			msg=f"""Max diff is {diff}"""					)





  def 					UpperCamelCase__   (      self:		Optional[Any]				,			UpperCamelCase_:		Optional[Any]				,			UpperCamelCase_:		str				,			UpperCamelCase_:		str					):
     UpperCamelCase_     =20
     UpperCamelCase_     =model_class_name(a_					)

     UpperCamelCase_    ,							UpperCamelCase_     =(
         inputs_dict["input_ids"],
         inputs_dict["attention_mask"],
     )

     UpperCamelCase_     =jnp.concatenate(
         [
             attention_mask,
             jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])					),
         ]				,			axis=-1				,			)

     UpperCamelCase_     =model.init_cache(input_ids.shape[0]				,			a_					)
     UpperCamelCase_     =jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1					)[None, :]				,			(input_ids.shape[0], input_ids.shape[-1] - 1)				,			)

     UpperCamelCase_     =model(
         input_ids[:, :-1]				,			attention_mask=a_				,			past_key_values=a_				,			position_ids=a_				,			)
     UpperCamelCase_     =jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]				,			dtype="i4"					)
     UpperCamelCase_     =model(
         input_ids[:, -1:]				,			past_key_values=outputs_cache.past_key_values				,			attention_mask=a_				,			position_ids=a_				,			)

     UpperCamelCase_     =model(a_				,			attention_mask=a_					)

     UpperCamelCase_     =np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])					)					)
     self.parent.assertTrue(diff < 1e-3				,			msg=f"""Max diff is {diff}"""					)




@require_flax
class 						__lowerCAmelCase   (			SCREAMING_SNAKE_CASE__       ,   unittest.TestCase       ,   SCREAMING_SNAKE_CASE__       ):







  '''simple docstring'''


  __lowerCamelCase       :    str     	=			(FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
  __lowerCamelCase       :    Optional[int]     	=			() if is_flax_available() else ()
  def 					UpperCamelCase__   (      self:		str					):
     UpperCamelCase_     =FlaxOPTModelTester(self					)
  def 					UpperCamelCase__   (      self:		Any					):
     UpperCamelCase_    ,							UpperCamelCase_     =self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
        self.model_tester.check_use_cache_forward(a_				,			a_				,			a_					)
  def 					UpperCamelCase__   (      self:		Optional[int]					):
     UpperCamelCase_    ,							UpperCamelCase_     =self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
        self.model_tester.check_use_cache_forward_with_attn_mask(a_				,			a_				,			a_					)





  @slow
  def 					UpperCamelCase__   (      self:		List[Any]					):



     for model_class_name in self.all_model_classes:
        UpperCamelCase_     =model_class_name.from_pretrained("facebook/opt-125m"					)
        UpperCamelCase_     =np.ones((1, 1)					) * model.config.eos_token_id
        UpperCamelCase_     =model(a_					)
        self.assertIsNotNone(a_					)




@require_sentencepiece
@require_flax
class 						__lowerCAmelCase   (			unittest.TestCase       ):







  '''simple docstring'''


  @slow
  def 					UpperCamelCase__   (      self:		Dict					):
     UpperCamelCase_     =FlaxOPTModel.from_pretrained("facebook/opt-350m"					)
     UpperCamelCase_     =jnp.array([[0, 3_1414, 232, 328, 740, 1140, 1_2695, 69, 4_6078, 1588, 2]]					)
     UpperCamelCase_     =model(input_ids=a_					).last_hidden_state
     UpperCamelCase_     =(1, 11, 512)
     self.assertEqual(output.shape				,			a_					)
     UpperCamelCase_     =jnp.array(
         [[-0.2867, -1.9256, -0.3062], [-1.2711, -0.1337, -0.1897], [0.4109, 0.1187, -1.3142]]					)
     self.assertTrue(jnp.allclose(output[:, :3, :3]				,			a_				,			atol=4e-2					)					)




@require_flax
@slow
class 						__lowerCAmelCase   (			unittest.TestCase       ):







  '''simple docstring'''


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





  def 					UpperCamelCase__   (      self:		str					):
     UpperCamelCase_     =FlaxOPTForCausalLM.from_pretrained(self.path_model					)
     UpperCamelCase_     =GPTaTokenizer.from_pretrained(self.path_model					)

     UpperCamelCase_     =[
         "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
     UpperCamelCase_     =tokenizer(a_				,			return_tensors="jax"				,			padding=a_				,			add_special_tokens=a_					)
     UpperCamelCase_     =model(inputs.input_ids				,			attention_mask=inputs.attention_mask					)[0].mean(axis=-1					)
     UpperCamelCase_     =jnp.array(
         [
             [1.3851, -13.8923, -10.5229, -10.7533, -0.2309, -10.2384, -0.5365, -9.0947, -5.1670],
             [-4.7073, -10.6276, -3.9415, -21.5242, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
             [0.6247, -3.4229, -8.9179, -1.4297, -14.1650, 1.4146, -9.0218, -0.2703, -0.2703],
             [6.4783, -1.9913, -10.7926, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
         ]					)
     self.assertTrue(jnp.allclose(a_				,			a_				,			atol=4e-2					)					)

     UpperCamelCase_     =jax.jit(a_					)
     UpperCamelCase_     =model(inputs.input_ids				,			attention_mask=inputs.attention_mask					)[0].mean(axis=-1					)
     self.assertTrue(jnp.allclose(a_				,			a_				,			atol=4e-2					)					)




@require_flax
@slow
class 						__lowerCAmelCase   (			unittest.TestCase       ):







  '''simple docstring'''


  @property
  def 					UpperCamelCase__   (      self:		Any					):
     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:		Union[str, Any]					):
     UpperCamelCase_     ="facebook/opt-125m"

     UpperCamelCase_     =[
         "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",
     ]

     UpperCamelCase_     =[]

     UpperCamelCase_     =FlaxOPTForCausalLM.from_pretrained(a_					)
     UpperCamelCase_     =GPTaTokenizer.from_pretrained(a_					)

     for prompt in self.prompts:
        UpperCamelCase_     =tokenizer(a_				,			return_tensors="jax"					).input_ids

        UpperCamelCase_     =model.generate(a_				,			max_length=10					)
        UpperCamelCase_     =generated_ids[0]

        UpperCamelCase_     =tokenizer.batch_decode(a_				,			skip_special_tokens=a_					)
        predicted_outputs += generated_string

     self.assertListEqual(a_				,			a_					)
  def 					UpperCamelCase__   (      self:		Any					):
     UpperCamelCase_     ="facebook/opt-350m"

     UpperCamelCase_     =[
         "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",
     ]

     UpperCamelCase_     =[]
     UpperCamelCase_     =FlaxOPTForCausalLM.from_pretrained(a_					)
     UpperCamelCase_     =GPTaTokenizer.from_pretrained(a_					)

     for prompt in self.prompts:
        UpperCamelCase_     =tokenizer(a_				,			return_tensors="jax"					).input_ids

        UpperCamelCase_     =model.generate(a_				,			max_length=10					)
        UpperCamelCase_     =generated_ids[0]

        UpperCamelCase_     =tokenizer.batch_decode(a_				,			skip_special_tokens=a_					)
        predicted_outputs += generated_string

     self.assertListEqual(a_				,			a_					)
  def 					UpperCamelCase__   (      self:		Optional[int]					):
     UpperCamelCase_     ="facebook/opt-125m"
     UpperCamelCase_     =[
         "Today is a beautiful day and I want to thank",
         "In the city of Rome Canaver Canaver Canaver Canaver",
     ]
     UpperCamelCase_     =FlaxOPTForCausalLM.from_pretrained(a_					)
     UpperCamelCase_     =GPTaTokenizer.from_pretrained(a_					)
     UpperCamelCase_     =tokenizer(
         [
             "Today is a beautiful day and I want to",
             "In the city of",
         ]				,			return_tensors="jax"				,			padding=a_				,			)

     UpperCamelCase_     =jax.jit(model.generate					)

     UpperCamelCase_     =jit_generate(inputs["input_ids"]				,			attention_mask=inputs["attention_mask"]					).sequences

     UpperCamelCase_     =tokenizer.batch_decode(a_				,			skip_special_tokens=a_					)

     self.assertIsNotNone(a_				,			a_					)





  def 					UpperCamelCase__   (      self:		Any					):
     UpperCamelCase_     ="facebook/opt-350m"

     UpperCamelCase_     =GPTaTokenizer.from_pretrained(a_					)
     UpperCamelCase_     =FlaxOPTForCausalLM.from_pretrained(a_					)

     UpperCamelCase_     ="left"

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

     UpperCamelCase_     =tokenizer(a_				,			return_tensors="jax"				,			padding=a_					)
     UpperCamelCase_     =inputs["input_ids"]

     UpperCamelCase_     =model.generate(input_ids=a_				,			attention_mask=inputs["attention_mask"]				,			trace=a_					)

     UpperCamelCase_     =tokenizer(sentences[0]				,			return_tensors="jax"					).input_ids
     UpperCamelCase_     =model.generate(input_ids=a_					)

     UpperCamelCase_     =inputs_non_padded.shape[-1] - inputs["attention_mask"][-1].sum()
     UpperCamelCase_     =tokenizer(sentences[1]				,			return_tensors="jax"					).input_ids
     UpperCamelCase_     =model.generate(input_ids=a_				,			max_length=model.config.max_length - num_paddings					)

     UpperCamelCase_     =tokenizer.batch_decode(outputs[0]				,			skip_special_tokens=a_					)
     UpperCamelCase_     =tokenizer.decode(output_non_padded[0][0]				,			skip_special_tokens=a_					)
     UpperCamelCase_     =tokenizer.decode(output_padded[0][0]				,			skip_special_tokens=a_					)

     UpperCamelCase_     =[
         "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]					)







"""simple docstring"""






from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFImageClassifierOutputWithNoAttention,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_resnet import ResNetConfig


A_			  =    logging.get_logger(__name__)

# General docstring
A_			  =    "ResNetConfig"

# Base docstring
A_			  =    "microsoft/resnet-50"
A_			  =    [1, 2048, 7, 7]

# Image classification docstring
A_			  =    "microsoft/resnet-50"
A_			  =    "tiger cat"

A_			  =    [
    "microsoft/resnet-50",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]
class 						__lowerCAmelCase   (			tf.keras.layers.Layer       ):







  '''simple docstring'''


  def __init__(      self:		Dict				,			UpperCamelCase_:		int				,			UpperCamelCase_:		int = 3				,			UpperCamelCase_:		int = 1				,			UpperCamelCase_:		str = "relu"				,			**UpperCamelCase_:		List[Any]					):
     super().__init__(**UpperCamelCase_					)
     UpperCamelCase_     =kernel_size // 2
     UpperCamelCase_     =tf.keras.layers.ConvaD(
         UpperCamelCase_				,			kernel_size=UpperCamelCase_				,			strides=UpperCamelCase_				,			padding="valid"				,			use_bias=UpperCamelCase_				,			name="convolution"					)
     # Use same default momentum and epsilon as PyTorch equivalent
     UpperCamelCase_     =tf.keras.layers.BatchNormalization(epsilon=1e-5				,			momentum=0.9				,			name="normalization"					)
     UpperCamelCase_     =ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("linear"					)
  def 					UpperCamelCase__   (      self:		List[str]				,			UpperCamelCase_:		tf.Tensor					):
     # Pad to match that done in the PyTorch Conv2D model
     UpperCamelCase_     =UpperCamelCase_     =(self.pad_value, self.pad_value)
     UpperCamelCase_     =tf.pad(UpperCamelCase_				,			[(0, 0), height_pad, width_pad, (0, 0)]					)
     UpperCamelCase_     =self.conv(UpperCamelCase_					)
     return hidden_state





  def 					UpperCamelCase__   (      self:		Optional[int]				,			UpperCamelCase_:		tf.Tensor				,			UpperCamelCase_:		bool = False					):
     UpperCamelCase_     =self.convolution(UpperCamelCase_					)
     UpperCamelCase_     =self.normalization(UpperCamelCase_				,			training=UpperCamelCase_					)
     UpperCamelCase_     =self.activation(UpperCamelCase_					)
     return hidden_state




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







  '''simple docstring'''


  def __init__(      self:		Optional[int]				,			UpperCamelCase_:		ResNetConfig				,			**UpperCamelCase_:		Dict					):
     super().__init__(**UpperCamelCase_					)
     UpperCamelCase_     =TFResNetConvLayer(
         config.embedding_size				,			kernel_size=7				,			stride=2				,			activation=config.hidden_act				,			name="embedder"				,			)
     UpperCamelCase_     =tf.keras.layers.MaxPoolaD(pool_size=3				,			strides=2				,			padding="valid"				,			name="pooler"					)
     UpperCamelCase_     =config.num_channels





  def 					UpperCamelCase__   (      self:		Dict				,			UpperCamelCase_:		tf.Tensor				,			UpperCamelCase_:		bool = False					):
     UpperCamelCase_    ,							UpperCamelCase_    ,							UpperCamelCase_    ,							UpperCamelCase_     =shape_list(UpperCamelCase_					)
     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."					)
     UpperCamelCase_     =pixel_values
     UpperCamelCase_     =self.embedder(UpperCamelCase_					)
     UpperCamelCase_     =tf.pad(UpperCamelCase_				,			[[0, 0], [1, 1], [1, 1], [0, 0]]					)
     UpperCamelCase_     =self.pooler(UpperCamelCase_					)
     return hidden_state




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







  '''simple docstring'''


  def __init__(      self:		Optional[int]				,			UpperCamelCase_:		int				,			UpperCamelCase_:		int = 2				,			**UpperCamelCase_:		Tuple					):
     super().__init__(**UpperCamelCase_					)
     UpperCamelCase_     =tf.keras.layers.ConvaD(
         UpperCamelCase_				,			kernel_size=1				,			strides=UpperCamelCase_				,			use_bias=UpperCamelCase_				,			name="convolution"					)
     # Use same default momentum and epsilon as PyTorch equivalent
     UpperCamelCase_     =tf.keras.layers.BatchNormalization(epsilon=1e-5				,			momentum=0.9				,			name="normalization"					)





  def 					UpperCamelCase__   (      self:		Optional[int]				,			UpperCamelCase_:		tf.Tensor				,			UpperCamelCase_:		bool = False					):
     UpperCamelCase_     =x
     UpperCamelCase_     =self.convolution(UpperCamelCase_					)
     UpperCamelCase_     =self.normalization(UpperCamelCase_				,			training=UpperCamelCase_					)
     return hidden_state




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







  '''simple docstring'''


  def __init__(      self:		Optional[Any]				,			UpperCamelCase_:		int				,			UpperCamelCase_:		int				,			UpperCamelCase_:		int = 1				,			UpperCamelCase_:		str = "relu"				,			**UpperCamelCase_:		Dict					):
     super().__init__(**UpperCamelCase_					)
     UpperCamelCase_     =in_channels != out_channels or stride != 1
     UpperCamelCase_     =TFResNetConvLayer(UpperCamelCase_				,			stride=UpperCamelCase_				,			name="layer.0"					)
     UpperCamelCase_     =TFResNetConvLayer(UpperCamelCase_				,			activation=UpperCamelCase_				,			name="layer.1"					)
     UpperCamelCase_     =(
         TFResNetShortCut(UpperCamelCase_				,			stride=UpperCamelCase_				,			name="shortcut"					)
         if should_apply_shortcut
         else tf.keras.layers.Activation("linear"				,			name="shortcut"					)
     )
     UpperCamelCase_     =ACTaFN[activation]





  def 					UpperCamelCase__   (      self:		Optional[Any]				,			UpperCamelCase_:		tf.Tensor				,			UpperCamelCase_:		bool = False					):
     UpperCamelCase_     =hidden_state
     UpperCamelCase_     =self.conva(UpperCamelCase_				,			training=UpperCamelCase_					)
     UpperCamelCase_     =self.conva(UpperCamelCase_				,			training=UpperCamelCase_					)
     UpperCamelCase_     =self.shortcut(UpperCamelCase_				,			training=UpperCamelCase_					)
     hidden_state += residual
     UpperCamelCase_     =self.activation(UpperCamelCase_					)
     return hidden_state




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







  '''simple docstring'''


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





  def 					UpperCamelCase__   (      self:		Optional[Any]				,			UpperCamelCase_:		tf.Tensor				,			UpperCamelCase_:		bool = False					):
     UpperCamelCase_     =hidden_state
     UpperCamelCase_     =self.conva(UpperCamelCase_				,			training=UpperCamelCase_					)
     UpperCamelCase_     =self.conva(UpperCamelCase_				,			training=UpperCamelCase_					)
     UpperCamelCase_     =self.conva(UpperCamelCase_				,			training=UpperCamelCase_					)
     UpperCamelCase_     =self.shortcut(UpperCamelCase_				,			training=UpperCamelCase_					)
     hidden_state += residual
     UpperCamelCase_     =self.activation(UpperCamelCase_					)
     return hidden_state




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







  '''simple docstring'''


  def __init__(      self:		Optional[int]				,			UpperCamelCase_:		ResNetConfig				,			UpperCamelCase_:		int				,			UpperCamelCase_:		int				,			UpperCamelCase_:		int = 2				,			UpperCamelCase_:		int = 2				,			**UpperCamelCase_:		Optional[Any]					):
     super().__init__(**UpperCamelCase_					)

     UpperCamelCase_     =TFResNetBottleNeckLayer if config.layer_type == "bottleneck" else TFResNetBasicLayer

     UpperCamelCase_     =[layer(UpperCamelCase_				,			UpperCamelCase_				,			stride=UpperCamelCase_				,			activation=config.hidden_act				,			name="layers.0"					)]
     layers += [
         layer(UpperCamelCase_				,			UpperCamelCase_				,			activation=config.hidden_act				,			name=f"""layers.{i + 1}"""					)
         for i in range(depth - 1					)
     ]
     UpperCamelCase_     =layers





  def 					UpperCamelCase__   (      self:		str				,			UpperCamelCase_:		tf.Tensor				,			UpperCamelCase_:		bool = False					):
     for layer in self.stage_layers:
        UpperCamelCase_     =layer(UpperCamelCase_				,			training=UpperCamelCase_					)
     return hidden_state




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







  '''simple docstring'''


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





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

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

        UpperCamelCase_     =stage_module(UpperCamelCase_				,			training=UpperCamelCase_					)

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




class 						__lowerCAmelCase   (			UpperCAmelCase       ):







  '''simple docstring'''


  __lowerCamelCase       :    Any     	=			ResNetConfig
  __lowerCamelCase       :    Dict     	=			"resnet"
  __lowerCamelCase       :    str     	=			"pixel_values"
  @property
  def 					UpperCamelCase__   (      self:		List[str]					):
     return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)				,			dtype=tf.floataa					)}




A_			  =    r"\n    This model is a TensorFlow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n"


A_			  =    r"\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"
@keras_serializable
class 						__lowerCAmelCase   (			tf.keras.layers.Layer       ):







  '''simple docstring'''


  __lowerCamelCase       :    str     	=			ResNetConfig
  def __init__(      self:		str				,			UpperCamelCase_:		ResNetConfig				,			**UpperCamelCase_:		Any					):
     super().__init__(**UpperCamelCase_					)
     UpperCamelCase_     =config
     UpperCamelCase_     =TFResNetEmbeddings(UpperCamelCase_				,			name="embedder"					)
     UpperCamelCase_     =TFResNetEncoder(UpperCamelCase_				,			name="encoder"					)
     UpperCamelCase_     =tf.keras.layers.GlobalAveragePoolingaD(keepdims=UpperCamelCase_					)





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

     # TF 2.0 image layers can't use NCHW format when running on CPU.
     # We transpose to NHWC format and then transpose back after the full forward pass.
     # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
     UpperCamelCase_     =tf.transpose(UpperCamelCase_				,			perm=[0, 2, 3, 1]					)
     UpperCamelCase_     =self.embedder(UpperCamelCase_				,			training=UpperCamelCase_					)

     UpperCamelCase_     =self.encoder(
         UpperCamelCase_				,			output_hidden_states=UpperCamelCase_				,			return_dict=UpperCamelCase_				,			training=UpperCamelCase_					)

     UpperCamelCase_     =encoder_outputs[0]

     UpperCamelCase_     =self.pooler(UpperCamelCase_					)

     # Transpose all the outputs to the NCHW format
     # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
     UpperCamelCase_     =tf.transpose(UpperCamelCase_				,			(0, 3, 1, 2)					)
     UpperCamelCase_     =tf.transpose(UpperCamelCase_				,			(0, 3, 1, 2)					)
     UpperCamelCase_     =()
     for hidden_state in encoder_outputs[1:]:
        UpperCamelCase_     =hidden_states + tuple(tf.transpose(UpperCamelCase_				,			(0, 3, 1, 2)					) for h in hidden_state					)

     if not return_dict:
        return (last_hidden_state, pooled_output) + hidden_states

     UpperCamelCase_     =hidden_states if output_hidden_states else None

     return TFBaseModelOutputWithPoolingAndNoAttention(
         last_hidden_state=UpperCamelCase_				,			pooler_output=UpperCamelCase_				,			hidden_states=UpperCamelCase_				,			)




@add_start_docstrings(
    "The bare ResNet model outputting raw features without any specific head on top."       ,   UpperCAmelCase       ,   )
class 						__lowerCAmelCase   (			UpperCAmelCase       ):







  '''simple docstring'''


  def __init__(      self:		Tuple				,			UpperCamelCase_:		ResNetConfig				,			**UpperCamelCase_:		int					):
     super().__init__(UpperCamelCase_				,			**UpperCamelCase_					)
     UpperCamelCase_     =TFResNetMainLayer(config=UpperCamelCase_				,			name="resnet"					)





  @add_start_docstrings_to_model_forward(UpperCamelCase_					)
  @add_code_sample_docstrings(
      checkpoint=_CHECKPOINT_FOR_DOC				,			output_type=UpperCamelCase_				,			config_class=_CONFIG_FOR_DOC				,			modality="vision"				,			expected_output=_EXPECTED_OUTPUT_SHAPE				,			)
  @unpack_inputs
  def 					UpperCamelCase__   (      self:		List[str]				,			UpperCamelCase_:		tf.Tensor				,			UpperCamelCase_:		Optional[bool] = None				,			UpperCamelCase_:		Optional[bool] = None				,			UpperCamelCase_:		bool = False				,			):
     UpperCamelCase_     =(
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
     UpperCamelCase_     =return_dict if return_dict is not None else self.config.use_return_dict

     UpperCamelCase_     =self.resnet(
         pixel_values=UpperCamelCase_				,			output_hidden_states=UpperCamelCase_				,			return_dict=UpperCamelCase_				,			training=UpperCamelCase_				,			)
     return resnet_outputs




@add_start_docstrings(
    "\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "       ,   UpperCAmelCase       ,   )
class 						__lowerCAmelCase   (			UpperCAmelCase       ,   UpperCAmelCase       ):







  '''simple docstring'''


  def __init__(      self:		List[Any]				,			UpperCamelCase_:		ResNetConfig				,			**UpperCamelCase_:		List[str]					):
     super().__init__(UpperCamelCase_				,			**UpperCamelCase_					)
     UpperCamelCase_     =config.num_labels
     UpperCamelCase_     =TFResNetMainLayer(UpperCamelCase_				,			name="resnet"					)
     # classification head
     UpperCamelCase_     =(
         tf.keras.layers.Dense(config.num_labels				,			name="classifier.1"					)
         if config.num_labels > 0
         else tf.keras.layers.Activation("linear"				,			name="classifier.1"					)
     )
  def 					UpperCamelCase__   (      self:		List[str]				,			UpperCamelCase_:		tf.Tensor					):
     UpperCamelCase_     =tf.keras.layers.Flatten()(UpperCamelCase_					)
     UpperCamelCase_     =self.classifier_layer(UpperCamelCase_					)
     return logits





  @add_start_docstrings_to_model_forward(UpperCamelCase_					)
  @add_code_sample_docstrings(
      checkpoint=_IMAGE_CLASS_CHECKPOINT				,			output_type=UpperCamelCase_				,			config_class=_CONFIG_FOR_DOC				,			expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT				,			)
  @unpack_inputs
  def 					UpperCamelCase__   (      self:		Any				,			UpperCamelCase_:		tf.Tensor = None				,			UpperCamelCase_:		tf.Tensor = None				,			UpperCamelCase_:		bool = None				,			UpperCamelCase_:		bool = None				,			UpperCamelCase_:		bool = False				,			):
     UpperCamelCase_     =return_dict if return_dict is not None else self.config.use_return_dict

     UpperCamelCase_     =self.resnet(
         UpperCamelCase_				,			output_hidden_states=UpperCamelCase_				,			return_dict=UpperCamelCase_				,			training=UpperCamelCase_					)

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

     UpperCamelCase_     =self.classifier(UpperCamelCase_					)

     UpperCamelCase_     =None if labels is None else self.hf_compute_loss(UpperCamelCase_				,			UpperCamelCase_					)

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

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