




import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

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


if is_flax_available():
  import jax
  import jax.numpy as jnp

  from transformers.modeling_flax_pytorch_utils import (
      convert_pytorch_state_dict_to_flax,
      load_flax_weights_in_pytorch_model,
  )
  from transformers.models.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
  import torch


class  __lowercase :







     '''simple docstring'''




     def __init__(			self    :     Tuple	,		UpperCamelCase_    :     Dict	,		UpperCamelCase_    :     Any=14	,		UpperCamelCase_    :     Dict=7	,		UpperCamelCase_    :     Optional[Any]=True	,		UpperCamelCase_    :     int=True	,		UpperCamelCase_    :     List[str]=False	,		UpperCamelCase_    :     Any=True	,		UpperCamelCase_    :     Union[str, Any]=99	,		UpperCamelCase_    :     Any=32	,		UpperCamelCase_    :     str=5	,		UpperCamelCase_    :     int=4	,		UpperCamelCase_    :     Tuple=37	,		UpperCamelCase_    :     Optional[int]="gelu"	,		UpperCamelCase_    :     str=0.1	,		UpperCamelCase_    :     Tuple=0.1	,		UpperCamelCase_    :     int=512	,		UpperCamelCase_    :     int=0.02	,		):


         """simple docstring"""



         __A             =		parent
         __A             =		batch_size
         __A             =		seq_length
         __A             =		is_training
         __A             =		use_input_mask
         __A             =		use_token_type_ids
         __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             =		initializer_range
         __A             =		None
         __A             =		vocab_size - 1
         __A             =		vocab_size - 1
         __A             =		vocab_size - 1




     def 						lowerCAmelCase_       (			self    :     List[Any]			):


         """simple docstring"""



         __A             =		ids_tensor([self.batch_size, self.seq_length]	,		self.vocab_size			)

         __A             =		None
         if self.use_input_mask:
             __A             =		random_attention_mask([self.batch_size, self.seq_length]			)

         __A             =		GPTaConfig(
             vocab_size=self.vocab_size	,		n_embd=self.hidden_size	,		n_layer=self.num_hidden_layers	,		n_head=self.num_attention_heads	,		n_positions=self.max_position_embeddings	,		use_cache=UpperCamelCase_	,		bos_token_id=self.bos_token_id	,		eos_token_id=self.eos_token_id	,		pad_token_id=self.pad_token_id	,		)

         return (config, input_ids, input_mask)




     def 						lowerCAmelCase_       (			self    :     Any			):


         """simple docstring"""



         __A             =		self.prepare_config_and_inputs()
         __A      ,		__A      ,		__A             =		config_and_inputs
         __A             =		{"""input_ids""": input_ids, """attention_mask""": attention_mask}
         return config, inputs_dict




     def 						lowerCAmelCase_       (			self    :     int			):


         """simple docstring"""



         __A      ,		__A      ,		__A             =		self.prepare_config_and_inputs()

         __A             =		floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
         __A             =		ids_tensor([self.batch_size, self.seq_length]	,		vocab_size=2			)

         return (
             config,
             input_ids,
             attention_mask,
             encoder_hidden_states,
             encoder_attention_mask,
         )




     def 						lowerCAmelCase_       (			self    :     Optional[int]	,		UpperCamelCase_    :     Any	,		UpperCamelCase_    :     Optional[int]	,		UpperCamelCase_    :     Tuple	,		UpperCamelCase_    :     str			):


         """simple docstring"""



         __A             =		20
         __A             =		model_class_name(UpperCamelCase_			)

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

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

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

         __A             =		model(UpperCamelCase_			)

         __A             =		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 						lowerCAmelCase_       (			self    :     Union[str, Any]	,		UpperCamelCase_    :     int	,		UpperCamelCase_    :     Optional[Any]	,		UpperCamelCase_    :     Union[str, Any]	,		UpperCamelCase_    :     List[str]			):


         """simple docstring"""



         __A             =		20
         __A             =		model_class_name(UpperCamelCase_			)

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

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

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

         __A             =		model(UpperCamelCase_	,		attention_mask=UpperCamelCase_			)

         __A             =		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  __lowercase (		lowercase_     , lowercase_     , unittest.TestCase					):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE				  =			(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
     SCREAMING_SNAKE_CASE				  =			(FlaxGPTaLMHeadModel,) if is_flax_available() else ()




     def 						lowerCAmelCase_       (			self    :     List[Any]			):


         """simple docstring"""



         __A             =		FlaxGPTaModelTester(self			)




     def 						lowerCAmelCase_       (			self    :     List[Any]			):


         """simple docstring"""



         for model_class_name in self.all_model_classes:
             __A      ,		__A      ,		__A             =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.check_use_cache_forward(UpperCamelCase_	,		UpperCamelCase_	,		UpperCamelCase_	,		UpperCamelCase_			)




     def 						lowerCAmelCase_       (			self    :     str			):


         """simple docstring"""



         for model_class_name in self.all_model_classes:
             __A      ,		__A      ,		__A             =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.check_use_cache_forward_with_attn_mask(
                 UpperCamelCase_	,		UpperCamelCase_	,		UpperCamelCase_	,		UpperCamelCase_			)




     @slow
     def 						lowerCAmelCase_       (			self    :     Optional[int]			):


         """simple docstring"""



         __A             =		GPTaTokenizer.from_pretrained("""gpt2"""	,		pad_token="""</s>"""	,		padding_side="""left"""			)
         __A             =		tokenizer(["""Hello this is a long string""", """Hey"""]	,		return_tensors="""np"""	,		padding=UpperCamelCase_	,		truncation=UpperCamelCase_			)

         __A             =		FlaxGPTaLMHeadModel.from_pretrained("""gpt2"""			)
         __A             =		False
         __A             =		model.config.eos_token_id

         __A             =		jax.jit(model.generate			)

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

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

         __A             =		[
             """Hello this is a long string of words. I\'m going to start with the first one.\n""",
             """Hey, I\'m not sure if I\'m going to be able to do""",
         ]

         self.assertListEqual(UpperCamelCase_	,		UpperCamelCase_			)




     @is_pt_flax_cross_test
     def 						lowerCAmelCase_       (			self    :     Optional[int]			):


         """simple docstring"""



         __A      ,		__A             =		self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__			):
                 # prepare inputs
                 __A             =		self._prepare_for_class(UpperCamelCase_	,		UpperCamelCase_			)
                 __A             =		{k: torch.tensor(v.tolist()			) for k, v in prepared_inputs_dict.items()}

                 # load corresponding PyTorch class
                 __A             =		model_class.__name__[4:]  # Skip the "Flax" at the beginning
                 __A             =		getattr(UpperCamelCase_	,		UpperCamelCase_			)

                 __A      ,		__A             =		pt_inputs["""input_ids"""].shape
                 __A             =		np.random.randint(0	,		seq_length - 1	,		size=(batch_size,)			)
                 for batch_idx, start_index in enumerate(UpperCamelCase_			):
                     __A             =		0
                     __A             =		1
                     __A             =		0
                     __A             =		1
                 __A             =		pt_model_class(UpperCamelCase_			).eval()
                 __A             =		model_class(UpperCamelCase_	,		dtype=jnp.floataa			)

                 __A             =		convert_pytorch_state_dict_to_flax(pt_model.state_dict()	,		UpperCamelCase_			)
                 __A             =		fx_state

                 with torch.no_grad():
                     __A             =		pt_model(**UpperCamelCase_			).to_tuple()

                 __A             =		fx_model(**UpperCamelCase_			).to_tuple()
                 self.assertEqual(len(UpperCamelCase_			)	,		len(UpperCamelCase_			)	,		"""Output lengths differ between Flax and PyTorch"""			)
                 for fx_output, pt_output in zip(UpperCamelCase_	,		UpperCamelCase_			):
                     self.assert_almost_equals(fx_output[:, -1]	,		pt_output[:, -1].numpy()	,		4e-2			)

                 with tempfile.TemporaryDirectory() as tmpdirname:
                     pt_model.save_pretrained(UpperCamelCase_			)
                     __A             =		model_class.from_pretrained(UpperCamelCase_	,		from_pt=UpperCamelCase_			)

                 __A             =		fx_model_loaded(**UpperCamelCase_			).to_tuple()
                 self.assertEqual(
                     len(UpperCamelCase_			)	,		len(UpperCamelCase_			)	,		"""Output lengths differ between Flax and PyTorch"""			)
                 for fx_output_loaded, pt_output in zip(UpperCamelCase_	,		UpperCamelCase_			):
                     self.assert_almost_equals(fx_output_loaded[:, -1]	,		pt_output[:, -1].numpy()	,		4e-2			)




     @is_pt_flax_cross_test
     def 						lowerCAmelCase_       (			self    :     int			):


         """simple docstring"""



         __A      ,		__A             =		self.model_tester.prepare_config_and_inputs_for_common()
         for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__			):
                 # prepare inputs
                 __A             =		self._prepare_for_class(UpperCamelCase_	,		UpperCamelCase_			)
                 __A             =		{k: torch.tensor(v.tolist()			) for k, v in prepared_inputs_dict.items()}

                 # load corresponding PyTorch class
                 __A             =		model_class.__name__[4:]  # Skip the "Flax" at the beginning
                 __A             =		getattr(UpperCamelCase_	,		UpperCamelCase_			)

                 __A             =		pt_model_class(UpperCamelCase_			).eval()
                 __A             =		model_class(UpperCamelCase_	,		dtype=jnp.floataa			)

                 __A             =		load_flax_weights_in_pytorch_model(UpperCamelCase_	,		fx_model.params			)
                 __A      ,		__A             =		pt_inputs["""input_ids"""].shape
                 __A             =		np.random.randint(0	,		seq_length - 1	,		size=(batch_size,)			)
                 for batch_idx, start_index in enumerate(UpperCamelCase_			):
                     __A             =		0
                     __A             =		1
                     __A             =		0
                     __A             =		1

                 # make sure weights are tied in PyTorch
                 pt_model.tie_weights()

                 with torch.no_grad():
                     __A             =		pt_model(**UpperCamelCase_			).to_tuple()

                 __A             =		fx_model(**UpperCamelCase_			).to_tuple()
                 self.assertEqual(len(UpperCamelCase_			)	,		len(UpperCamelCase_			)	,		"""Output lengths differ between Flax and PyTorch"""			)
                 for fx_output, pt_output in zip(UpperCamelCase_	,		UpperCamelCase_			):
                     self.assert_almost_equals(fx_output[:, -1]	,		pt_output[:, -1].numpy()	,		4e-2			)

                 with tempfile.TemporaryDirectory() as tmpdirname:
                     fx_model.save_pretrained(UpperCamelCase_			)
                     __A             =		pt_model_class.from_pretrained(UpperCamelCase_	,		from_flax=UpperCamelCase_			)

                 with torch.no_grad():
                     __A             =		pt_model_loaded(**UpperCamelCase_			).to_tuple()

                 self.assertEqual(
                     len(UpperCamelCase_			)	,		len(UpperCamelCase_			)	,		"""Output lengths differ between Flax and PyTorch"""			)
                 for fx_output, pt_output in zip(UpperCamelCase_	,		UpperCamelCase_			):
                     self.assert_almost_equals(fx_output[:, -1]	,		pt_output[:, -1].numpy()	,		4e-2			)





     @slow
     def 						lowerCAmelCase_       (			self    :     Any			):


         """simple docstring"""



         for model_class_name in self.all_model_classes:
             __A             =		model_class_name.from_pretrained("""gpt2"""	,		from_pt=UpperCamelCase_			)
             __A             =		model(np.ones((1, 1)			)			)
             self.assertIsNotNone(UpperCamelCase_			)












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 :    List[Any]								=				logging.get_logger(__name__)

# General docstring
__a :    Any								=				"ResNetConfig"

# Base docstring
__a :    List[str]								=				"microsoft/resnet-50"
__a :    Dict								=				[1, 2048, 7, 7]

# Image classification docstring
__a :    Optional[int]								=				"microsoft/resnet-50"
__a :    Dict								=				"tiger cat"

__a :    Dict								=				[
    "microsoft/resnet-50",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]


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







     '''simple docstring'''




     def __init__(			self    :     Dict	,		UpperCamelCase_    :     int	,		UpperCamelCase_    :     int = 3	,		UpperCamelCase_    :     int = 1	,		UpperCamelCase_    :     str = "relu"	,		**UpperCamelCase_    :     Dict			):


         """simple docstring"""



         super().__init__(**UpperCamelCase_			)
         __A             =		kernel_size // 2
         __A             =		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
         __A             =		tf.keras.layers.BatchNormalization(epsilon=1e-5	,		momentum=0.9	,		name="""normalization"""			)
         __A             =		ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("""linear"""			)




     def 						lowerCAmelCase_       (			self    :     Any	,		UpperCamelCase_    :     tf.Tensor			):


         """simple docstring"""



         __A             =		__A             =		(self.pad_value, self.pad_value)
         __A             =		tf.pad(UpperCamelCase_	,		[(0, 0), height_pad, width_pad, (0, 0)]			)
         __A             =		self.conv(UpperCamelCase_			)
         return hidden_state





     def 						lowerCAmelCase_       (			self    :     Optional[int]	,		UpperCamelCase_    :     tf.Tensor	,		UpperCamelCase_    :     bool = False			):


         """simple docstring"""



         __A             =		self.convolution(UpperCamelCase_			)
         __A             =		self.normalization(UpperCamelCase_	,		training=UpperCamelCase_			)
         __A             =		self.activation(UpperCamelCase_			)
         return hidden_state


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







     '''simple docstring'''




     def __init__(			self    :     int	,		UpperCamelCase_    :     ResNetConfig	,		**UpperCamelCase_    :     Tuple			):


         """simple docstring"""



         super().__init__(**UpperCamelCase_			)
         __A             =		TFResNetConvLayer(
             config.embedding_size	,		kernel_size=7	,		stride=2	,		activation=config.hidden_act	,		name="""embedder"""	,		)
         __A             =		tf.keras.layers.MaxPoolaD(pool_size=3	,		strides=2	,		padding="""valid"""	,		name="""pooler"""			)
         __A             =		config.num_channels





     def 						lowerCAmelCase_       (			self    :     Any	,		UpperCamelCase_    :     tf.Tensor	,		UpperCamelCase_    :     bool = False			):


         """simple docstring"""



         __A      ,		__A      ,		__A      ,		__A             =		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."""			)
         __A             =		pixel_values
         __A             =		self.embedder(UpperCamelCase_			)
         __A             =		tf.pad(UpperCamelCase_	,		[[0, 0], [1, 1], [1, 1], [0, 0]]			)
         __A             =		self.pooler(UpperCamelCase_			)
         return hidden_state


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







     '''simple docstring'''




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


         """simple docstring"""



         super().__init__(**UpperCamelCase_			)
         __A             =		tf.keras.layers.ConvaD(
             UpperCamelCase_	,		kernel_size=1	,		strides=UpperCamelCase_	,		use_bias=UpperCamelCase_	,		name="""convolution"""			)
         # Use same default momentum and epsilon as PyTorch equivalent
         __A             =		tf.keras.layers.BatchNormalization(epsilon=1e-5	,		momentum=0.9	,		name="""normalization"""			)





     def 						lowerCAmelCase_       (			self    :     str	,		UpperCamelCase_    :     tf.Tensor	,		UpperCamelCase_    :     bool = False			):


         """simple docstring"""



         __A             =		x
         __A             =		self.convolution(UpperCamelCase_			)
         __A             =		self.normalization(UpperCamelCase_	,		training=UpperCamelCase_			)
         return hidden_state


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







     '''simple docstring'''




     def __init__(			self    :     List[str]	,		UpperCamelCase_    :     int	,		UpperCamelCase_    :     int	,		UpperCamelCase_    :     int = 1	,		UpperCamelCase_    :     str = "relu"	,		**UpperCamelCase_    :     Any			):


         """simple docstring"""



         super().__init__(**UpperCamelCase_			)
         __A             =		in_channels != out_channels or stride != 1
         __A             =		TFResNetConvLayer(UpperCamelCase_	,		stride=UpperCamelCase_	,		name="""layer.0"""			)
         __A             =		TFResNetConvLayer(UpperCamelCase_	,		activation=UpperCamelCase_	,		name="""layer.1"""			)
         __A             =		(
             TFResNetShortCut(UpperCamelCase_	,		stride=UpperCamelCase_	,		name="""shortcut"""			)
             if should_apply_shortcut
             else tf.keras.layers.Activation("""linear"""	,		name="""shortcut"""			)
         )
         __A             =		ACTaFN[activation]





     def 						lowerCAmelCase_       (			self    :     Union[str, Any]	,		UpperCamelCase_    :     tf.Tensor	,		UpperCamelCase_    :     bool = False			):


         """simple docstring"""



         __A             =		hidden_state
         __A             =		self.conva(UpperCamelCase_	,		training=UpperCamelCase_			)
         __A             =		self.conva(UpperCamelCase_	,		training=UpperCamelCase_			)
         __A             =		self.shortcut(UpperCamelCase_	,		training=UpperCamelCase_			)
         hidden_state += residual
         __A             =		self.activation(UpperCamelCase_			)
         return hidden_state


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







     '''simple docstring'''




     def __init__(			self    :     Optional[Any]	,		UpperCamelCase_    :     int	,		UpperCamelCase_    :     int	,		UpperCamelCase_    :     int = 1	,		UpperCamelCase_    :     str = "relu"	,		UpperCamelCase_    :     int = 4	,		**UpperCamelCase_    :     str	,		):


         """simple docstring"""



         super().__init__(**UpperCamelCase_			)
         __A             =		in_channels != out_channels or stride != 1
         __A             =		out_channels // reduction
         __A             =		TFResNetConvLayer(UpperCamelCase_	,		kernel_size=1	,		name="""layer.0"""			)
         __A             =		TFResNetConvLayer(UpperCamelCase_	,		stride=UpperCamelCase_	,		name="""layer.1"""			)
         __A             =		TFResNetConvLayer(UpperCamelCase_	,		kernel_size=1	,		activation=UpperCamelCase_	,		name="""layer.2"""			)
         __A             =		(
             TFResNetShortCut(UpperCamelCase_	,		stride=UpperCamelCase_	,		name="""shortcut"""			)
             if should_apply_shortcut
             else tf.keras.layers.Activation("""linear"""	,		name="""shortcut"""			)
         )
         __A             =		ACTaFN[activation]





     def 						lowerCAmelCase_       (			self    :     Any	,		UpperCamelCase_    :     tf.Tensor	,		UpperCamelCase_    :     bool = False			):


         """simple docstring"""



         __A             =		hidden_state
         __A             =		self.conva(UpperCamelCase_	,		training=UpperCamelCase_			)
         __A             =		self.conva(UpperCamelCase_	,		training=UpperCamelCase_			)
         __A             =		self.conva(UpperCamelCase_	,		training=UpperCamelCase_			)
         __A             =		self.shortcut(UpperCamelCase_	,		training=UpperCamelCase_			)
         hidden_state += residual
         __A             =		self.activation(UpperCamelCase_			)
         return hidden_state


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







     '''simple docstring'''




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


         """simple docstring"""



         super().__init__(**UpperCamelCase_			)

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

         __A             =		[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			)
         ]
         __A             =		layers





     def 						lowerCAmelCase_       (			self    :     List[str]	,		UpperCamelCase_    :     tf.Tensor	,		UpperCamelCase_    :     bool = False			):


         """simple docstring"""



         for layer in self.stage_layers:
             __A             =		layer(UpperCamelCase_	,		training=UpperCamelCase_			)
         return hidden_state


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







     '''simple docstring'''




     def __init__(			self    :     int	,		UpperCamelCase_    :     ResNetConfig	,		**UpperCamelCase_    :     Dict			):


         """simple docstring"""



         super().__init__(**UpperCamelCase_			)
         # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
         __A             =		[
             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 						lowerCAmelCase_       (			self    :     Optional[Any]	,		UpperCamelCase_    :     tf.Tensor	,		UpperCamelCase_    :     bool = False	,		UpperCamelCase_    :     bool = True	,		UpperCamelCase_    :     bool = False	,		):


         """simple docstring"""



         __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(UpperCamelCase_	,		training=UpperCamelCase_			)

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


class  __lowercase (		lowercase_					):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE				  =			ResNetConfig
     SCREAMING_SNAKE_CASE				  =			"resnet"
     SCREAMING_SNAKE_CASE				  =			"pixel_values"




     @property
     def 						lowerCAmelCase_       (			self    :     str			):


         """simple docstring"""



         return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)	,		dtype=tf.floataa			)}
__a :    Union[str, Any]								=				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 :    Tuple								=				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  __lowercase (		tf.keras.layers.Layer					):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE				  =			ResNetConfig




     def __init__(			self    :     str	,		UpperCamelCase_    :     ResNetConfig	,		**UpperCamelCase_    :     str			):


         """simple docstring"""



         super().__init__(**UpperCamelCase_			)
         __A             =		config
         __A             =		TFResNetEmbeddings(UpperCamelCase_	,		name="""embedder"""			)
         __A             =		TFResNetEncoder(UpperCamelCase_	,		name="""encoder"""			)
         __A             =		tf.keras.layers.GlobalAveragePoolingaD(keepdims=UpperCamelCase_			)





     @unpack_inputs
     def 						lowerCAmelCase_       (			self    :     str	,		UpperCamelCase_    :     tf.Tensor	,		UpperCamelCase_    :     Optional[bool] = None	,		UpperCamelCase_    :     Optional[bool] = None	,		UpperCamelCase_    :     bool = False	,		):


         """simple docstring"""



         __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

         # 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)
         __A             =		tf.transpose(UpperCamelCase_	,		perm=[0, 2, 3, 1]			)
         __A             =		self.embedder(UpperCamelCase_	,		training=UpperCamelCase_			)

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

         __A             =		encoder_outputs[0]

         __A             =		self.pooler(UpperCamelCase_			)

         # Transpose all the outputs to the NCHW format
         # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
         __A             =		tf.transpose(UpperCamelCase_	,		(0, 3, 1, 2)			)
         __A             =		tf.transpose(UpperCamelCase_	,		(0, 3, 1, 2)			)
         __A             =		()
         for hidden_state in encoder_outputs[1:]:
             __A             =		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

         __A             =		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."     , lowercase_     , )
class  __lowercase (		lowercase_					):







     '''simple docstring'''




     def __init__(			self    :     int	,		UpperCamelCase_    :     ResNetConfig	,		**UpperCamelCase_    :     Dict			):


         """simple docstring"""



         super().__init__(UpperCamelCase_	,		**UpperCamelCase_			)
         __A             =		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 						lowerCAmelCase_       (			self    :     Dict	,		UpperCamelCase_    :     tf.Tensor	,		UpperCamelCase_    :     Optional[bool] = None	,		UpperCamelCase_    :     Optional[bool] = None	,		UpperCamelCase_    :     bool = False	,		):


         """simple docstring"""



         __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.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    "     , lowercase_     , )
class  __lowercase (		lowercase_     , lowercase_					):







     '''simple docstring'''




     def __init__(			self    :     int	,		UpperCamelCase_    :     ResNetConfig	,		**UpperCamelCase_    :     Any			):


         """simple docstring"""



         super().__init__(UpperCamelCase_	,		**UpperCamelCase_			)
         __A             =		config.num_labels
         __A             =		TFResNetMainLayer(UpperCamelCase_	,		name="""resnet"""			)
         # classification head
         __A             =		(
             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 						lowerCAmelCase_       (			self    :     Union[str, Any]	,		UpperCamelCase_    :     tf.Tensor			):


         """simple docstring"""



         __A             =		tf.keras.layers.Flatten()(UpperCamelCase_			)
         __A             =		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 						lowerCAmelCase_       (			self    :     str	,		UpperCamelCase_    :     tf.Tensor = None	,		UpperCamelCase_    :     tf.Tensor = None	,		UpperCamelCase_    :     bool = None	,		UpperCamelCase_    :     bool = None	,		UpperCamelCase_    :     bool = False	,		):


         """simple docstring"""



         __A             =		return_dict if return_dict is not None else self.config.use_return_dict

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

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

         __A             =		self.classifier(UpperCamelCase_			)

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

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

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




