



import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import TimesformerConfig, TimesformerForVideoClassification, VideoMAEImageProcessor



def  lowercase    (				SCREAMING_SNAKE_CASE__    :		str			)					->       Optional[int]:






 """simple docstring"""




 _snake_case      :					Optional[int]					=   TimesformerConfig()

 if "large" in model_name:
  _snake_case      :					str					=   96

 if "hr" in model_name:
  _snake_case      :					str					=   16
  _snake_case      :					Optional[Any]					=   448

 _snake_case      :					List[str]					=   """huggingface/label-files"""
 if "k400" in model_name:
  _snake_case      :					Optional[int]					=   400
  _snake_case      :					Tuple					=   """kinetics400-id2label.json"""
 elif "k600" in model_name:
  _snake_case      :					Optional[Any]					=   600
  _snake_case      :					int					=   """kinetics600-id2label.json"""
 elif "ssv2" in model_name:
  _snake_case      :					Union[str, Any]					=   174
  _snake_case      :					List[str]					=   """something-something-v2-id2label.json"""
 else:
  raise ValueError("""Model name should either contain \'k400\', \'k600\' or \'ssv2\'."""			)
 _snake_case      :					Any					=   json.load(open(hf_hub_download(lowercase_			,    lowercase_			,    repo_type="""dataset"""			)			,    """r"""			)			)
 _snake_case      :					Dict					=   {int(lowercase_			): v for k, v in idalabel.items()}
 _snake_case      :					str					=   idalabel
 _snake_case      :					Dict					=   {v: k for k, v in idalabel.items()}

 return config



def  lowercase    (				SCREAMING_SNAKE_CASE__    :		List[str]			)					->       Any:






 """simple docstring"""




 if "encoder." in name:
  _snake_case      :					Union[str, Any]					=   name.replace("""encoder."""			,    """"""			)
 if "cls_token" in name:
  _snake_case      :					Optional[int]					=   name.replace("""cls_token"""			,    """timesformer.embeddings.cls_token"""			)
 if "pos_embed" in name:
  _snake_case      :					Dict					=   name.replace("""pos_embed"""			,    """timesformer.embeddings.position_embeddings"""			)
 if "time_embed" in name:
  _snake_case      :					Dict					=   name.replace("""time_embed"""			,    """timesformer.embeddings.time_embeddings"""			)
 if "patch_embed.proj" in name:
  _snake_case      :					Union[str, Any]					=   name.replace("""patch_embed.proj"""			,    """timesformer.embeddings.patch_embeddings.projection"""			)
 if "patch_embed.norm" in name:
  _snake_case      :					Tuple					=   name.replace("""patch_embed.norm"""			,    """timesformer.embeddings.norm"""			)
 if "blocks" in name:
  _snake_case      :					Dict					=   name.replace("""blocks"""			,    """timesformer.encoder.layer"""			)
 if "attn.proj" in name:
  _snake_case      :					Any					=   name.replace("""attn.proj"""			,    """attention.output.dense"""			)
 if "attn" in name and "bias" not in name and "temporal" not in name:
  _snake_case      :					Dict					=   name.replace("""attn"""			,    """attention.self"""			)
 if "attn" in name and "temporal" not in name:
  _snake_case      :					Any					=   name.replace("""attn"""			,    """attention.attention"""			)
 if "temporal_norm1" in name:
  _snake_case      :					str					=   name.replace("""temporal_norm1"""			,    """temporal_layernorm"""			)
 if "temporal_attn.proj" in name:
  _snake_case      :					List[str]					=   name.replace("""temporal_attn"""			,    """temporal_attention.output.dense"""			)
 if "temporal_fc" in name:
  _snake_case      :					Any					=   name.replace("""temporal_fc"""			,    """temporal_dense"""			)
 if "norm1" in name and "temporal" not in name:
  _snake_case      :					str					=   name.replace("""norm1"""			,    """layernorm_before"""			)
 if "norm2" in name:
  _snake_case      :					Tuple					=   name.replace("""norm2"""			,    """layernorm_after"""			)
 if "mlp.fc1" in name:
  _snake_case      :					Optional[int]					=   name.replace("""mlp.fc1"""			,    """intermediate.dense"""			)
 if "mlp.fc2" in name:
  _snake_case      :					int					=   name.replace("""mlp.fc2"""			,    """output.dense"""			)
 if "norm.weight" in name and "fc" not in name and "temporal" not in name:
  _snake_case      :					List[Any]					=   name.replace("""norm.weight"""			,    """timesformer.layernorm.weight"""			)
 if "norm.bias" in name and "fc" not in name and "temporal" not in name:
  _snake_case      :					Optional[int]					=   name.replace("""norm.bias"""			,    """timesformer.layernorm.bias"""			)
 if "head" in name:
  _snake_case      :					List[str]					=   name.replace("""head"""			,    """classifier"""			)

 return name



def  lowercase    (				SCREAMING_SNAKE_CASE__    :		List[Any]			,    SCREAMING_SNAKE_CASE__    :		int			)					->       List[Any]:






 """simple docstring"""




 for key in orig_state_dict.copy().keys():
  _snake_case      :					Any					=   orig_state_dict.pop(lowercase_			)

  if key.startswith("""model."""			):
   _snake_case      :					Optional[int]					=   key.replace("""model."""			,    """"""			)

  if "qkv" in key:
   _snake_case      :					Dict					=   key.split("""."""			)
   _snake_case      :					Union[str, Any]					=   int(key_split[1]			)
   _snake_case      :					Union[str, Any]					=   """timesformer.encoder.layer."""
   if "temporal" in key:
    _snake_case      :					Any					=   """.temporal_attention.attention.qkv."""
   else:
    _snake_case      :					str					=   """.attention.attention.qkv."""
   if "weight" in key:
    _snake_case      :					Optional[int]					=   val
   else:
    _snake_case      :					int					=   val
  else:
   _snake_case      :					Tuple					=   val

 return orig_state_dict



def  lowercase    (				)					->       Any:






 """simple docstring"""




 _snake_case      :					Optional[Any]					=   hf_hub_download(
     repo_id="""hf-internal-testing/spaghetti-video"""			,    filename="""eating_spaghetti.npy"""			,    repo_type="""dataset"""			)
 _snake_case      :					Any					=   np.load(lowercase_			)
 return list(lowercase_			)



def  lowercase    (				SCREAMING_SNAKE_CASE__    :		Optional[int]			,    SCREAMING_SNAKE_CASE__    :		Dict			,    SCREAMING_SNAKE_CASE__    :		List[Any]			,    SCREAMING_SNAKE_CASE__    :		Tuple			)					->       List[Any]:






 """simple docstring"""




 _snake_case      :					str					=   get_timesformer_config(lowercase_			)

 _snake_case      :					Any					=   TimesformerForVideoClassification(lowercase_			)

 # download original checkpoint, hosted on Google Drive
 _snake_case      :					Dict					=   """pytorch_model.bin"""
 gdown.cached_download(lowercase_			,    lowercase_			,    quiet=lowercase_			)
 _snake_case      :					Dict					=   torch.load(lowercase_			,    map_location="""cpu"""			)
 if "model" in files:
  _snake_case      :					Optional[Any]					=   files["""model"""]
 elif "module" in files:
  _snake_case      :					Union[str, Any]					=   files["""module"""]
 else:
  _snake_case      :					Any					=   files["""model_state"""]
 _snake_case      :					Any					=   convert_state_dict(lowercase_			,    lowercase_			)

 model.load_state_dict(lowercase_			)
 model.eval()

 # verify model on basic input
 _snake_case      :					int					=   VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]			,    image_std=[0.5, 0.5, 0.5]			)
 _snake_case      :					Dict					=   prepare_video()
 _snake_case      :					int					=   image_processor(video[:8]			,    return_tensors="""pt"""			)

 _snake_case      :					Optional[int]					=   model(**lowercase_			)
 _snake_case      :					Optional[int]					=   outputs.logits

 _snake_case      :					Optional[int]					=   [
     # Kinetics-400 checkpoints (hr = high resolution input of 448px instead of 224px)
     """timesformer-base-finetuned-k400""",
     """timesformer-large-finetuned-k400""",
     """timesformer-hr-finetuned-k400""",
     # Kinetics-600 checkpoints (hr = high resolution input of 448px instead of 224px)
     """timesformer-base-finetuned-k600""",
     """timesformer-large-finetuned-k600""",
     """timesformer-hr-finetuned-k600""",
     # Something-Something-v2 checkpoints (hr = high resolution input of 448px instead of 224px)
     """timesformer-base-finetuned-ssv2""",
     """timesformer-large-finetuned-ssv2""",
     """timesformer-hr-finetuned-ssv2""",
 ]

 # NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
 if model_name == "timesformer-base-finetuned-k400":
  _snake_case      :					str					=   torch.Size([1, 400]			)
  _snake_case      :					Optional[int]					=   torch.tensor([-0.3_0_1_6, -0.7_7_1_3, -0.4_2_0_5]			)
 elif model_name == "timesformer-base-finetuned-k600":
  _snake_case      :					Tuple					=   torch.Size([1, 600]			)
  _snake_case      :					Optional[Any]					=   torch.tensor([-0.7_2_6_7, -0.7_4_6_6, 3.2_4_0_4]			)
 elif model_name == "timesformer-base-finetuned-ssv2":
  _snake_case      :					Tuple					=   torch.Size([1, 174]			)
  _snake_case      :					List[str]					=   torch.tensor([-0.9_0_5_9, 0.6_4_3_3, -3.1_4_5_7]			)
 elif model_name == "timesformer-large-finetuned-k400":
  _snake_case      :					Optional[int]					=   torch.Size([1, 400]			)
  _snake_case      :					Union[str, Any]					=   torch.tensor([0, 0, 0]			)
 elif model_name == "timesformer-large-finetuned-k600":
  _snake_case      :					Union[str, Any]					=   torch.Size([1, 600]			)
  _snake_case      :					List[str]					=   torch.tensor([0, 0, 0]			)
 elif model_name == "timesformer-large-finetuned-ssv2":
  _snake_case      :					Tuple					=   torch.Size([1, 174]			)
  _snake_case      :					int					=   torch.tensor([0, 0, 0]			)
 elif model_name == "timesformer-hr-finetuned-k400":
  _snake_case      :					Tuple					=   torch.Size([1, 400]			)
  _snake_case      :					List[Any]					=   torch.tensor([-0.9_6_1_7, -3.7_3_1_1, -3.7_7_0_8]			)
 elif model_name == "timesformer-hr-finetuned-k600":
  _snake_case      :					Optional[Any]					=   torch.Size([1, 600]			)
  _snake_case      :					Optional[int]					=   torch.tensor([2.5_2_7_3, 0.7_1_2_7, 1.8_8_4_8]			)
 elif model_name == "timesformer-hr-finetuned-ssv2":
  _snake_case      :					int					=   torch.Size([1, 174]			)
  _snake_case      :					List[str]					=   torch.tensor([-3.6_7_5_6, -0.7_5_1_3, 0.7_1_8_0]			)
 else:
  raise ValueError(F'''Model name not supported. Should be one of {model_names}'''			)

 # verify logits
 assert logits.shape == expected_shape
 assert torch.allclose(logits[0, :3]			,    lowercase_			,    atol=1e-4			)
 print("""Logits ok!"""			)

 if pytorch_dump_folder_path is not None:
  print(F'''Saving model and image processor to {pytorch_dump_folder_path}'''			)
  image_processor.save_pretrained(lowercase_			)
  model.save_pretrained(lowercase_			)

 if push_to_hub:
  print("""Pushing to the hub..."""			)
  model.push_to_hub(F'''fcakyon/{model_name}'''			)


if __name__ == "__main__":
  a__      =		argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--checkpoint_url""",
      default="""https://drive.google.com/u/1/uc?id=17yvuYp9L4mn-HpIcK5Zo6K3UoOy1kA5l&export=download""",
      type=str,
      help=(
          """URL of the original PyTorch checkpoint (on Google Drive) you'd like to convert. Should be a direct"""
          """ download link."""
      ),
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""",
      default="""""",
      type=str,
      help="""Path to the output PyTorch model directory.""",
  )
  parser.add_argument("""--model_name""", default="""timesformer-base-finetuned-k400""", type=str, help="""Name of the model.""")
  parser.add_argument(
      """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
  )

  a__      =		parser.parse_args()
  convert_timesformer_checkpoint(
      args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub
  )





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, 20_48, 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 						snake_case  (		tf.keras.layers.Layer					):



   '''simple docstring'''






   def __init__(   self							:     Union[str, Any]			,						lowerCAmelCase							:     int			,						lowerCAmelCase							:     int = 3			,						lowerCAmelCase							:     int = 1			,						lowerCAmelCase							:     str = "relu"			,						**lowerCAmelCase							:     Any)			->			None:

    """simple docstring"""







    super().__init__(**lowerCAmelCase)
    _snake_case      :					Optional[Any]					=   kernel_size // 2
    _snake_case      :					int					=   tf.keras.layers.ConvaD(
        lowerCAmelCase			,						kernel_size=lowerCAmelCase			,						strides=lowerCAmelCase			,						padding="""valid"""			,						use_bias=lowerCAmelCase			,						name="""convolution""")
    # Use same default momentum and epsilon as PyTorch equivalent
    _snake_case      :					Optional[int]					=   tf.keras.layers.BatchNormalization(epsilon=1E-5			,						momentum=0.9			,						name="""normalization""")
    _snake_case      :					Union[str, Any]					=   ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("""linear""")

   def        UpperCamelCase_					(   self							:     Tuple			,						lowerCAmelCase							:     tf.Tensor)			->			tf.Tensor:

    """simple docstring"""







    _snake_case      :					List[Any]					=   (self.pad_value, self.pad_value)
    _snake_case      :					Optional[Any]					=   tf.pad(lowerCAmelCase			,						[(0, 0), height_pad, width_pad, (0, 0)])
    _snake_case      :					Optional[Any]					=   self.conv(lowerCAmelCase)
    return hidden_state


   def        UpperCamelCase_					(   self							:     Optional[Any]			,						lowerCAmelCase							:     tf.Tensor			,						lowerCAmelCase							:     bool = False)			->			tf.Tensor:

    """simple docstring"""







    _snake_case      :					Optional[Any]					=   self.convolution(lowerCAmelCase)
    _snake_case      :					List[Any]					=   self.normalization(lowerCAmelCase			,						training=lowerCAmelCase)
    _snake_case      :					List[Any]					=   self.activation(lowerCAmelCase)
    return hidden_state







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



   '''simple docstring'''






   def __init__(   self							:     Optional[int]			,						lowerCAmelCase							:     ResNetConfig			,						**lowerCAmelCase							:     Optional[Any])			->			None:

    """simple docstring"""







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


   def        UpperCamelCase_					(   self							:     List[Any]			,						lowerCAmelCase							:     tf.Tensor			,						lowerCAmelCase							:     bool = False)			->			tf.Tensor:

    """simple docstring"""







    _snake_case		,      _snake_case		,      _snake_case		,      _snake_case      :					Optional[Any]					=   shape_list(lowerCAmelCase)
    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.""")
    _snake_case      :					str					=   pixel_values
    _snake_case      :					str					=   self.embedder(lowerCAmelCase)
    _snake_case      :					Tuple					=   tf.pad(lowerCAmelCase			,						[[0, 0], [1, 1], [1, 1], [0, 0]])
    _snake_case      :					Dict					=   self.pooler(lowerCAmelCase)
    return hidden_state







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



   '''simple docstring'''






   def __init__(   self							:     Union[str, Any]			,						lowerCAmelCase							:     int			,						lowerCAmelCase							:     int = 2			,						**lowerCAmelCase							:     str)			->			None:

    """simple docstring"""







    super().__init__(**lowerCAmelCase)
    _snake_case      :					Any					=   tf.keras.layers.ConvaD(
        lowerCAmelCase			,						kernel_size=1			,						strides=lowerCAmelCase			,						use_bias=lowerCAmelCase			,						name="""convolution""")
    # Use same default momentum and epsilon as PyTorch equivalent
    _snake_case      :					Tuple					=   tf.keras.layers.BatchNormalization(epsilon=1E-5			,						momentum=0.9			,						name="""normalization""")


   def        UpperCamelCase_					(   self							:     List[Any]			,						lowerCAmelCase							:     tf.Tensor			,						lowerCAmelCase							:     bool = False)			->			tf.Tensor:

    """simple docstring"""







    _snake_case      :					Optional[int]					=   x
    _snake_case      :					Optional[Any]					=   self.convolution(lowerCAmelCase)
    _snake_case      :					int					=   self.normalization(lowerCAmelCase			,						training=lowerCAmelCase)
    return hidden_state







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



   '''simple docstring'''






   def __init__(   self							:     List[str]			,						lowerCAmelCase							:     int			,						lowerCAmelCase							:     int			,						lowerCAmelCase							:     int = 1			,						lowerCAmelCase							:     str = "relu"			,						**lowerCAmelCase							:     Tuple)			->			None:

    """simple docstring"""







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


   def        UpperCamelCase_					(   self							:     Optional[Any]			,						lowerCAmelCase							:     tf.Tensor			,						lowerCAmelCase							:     bool = False)			->			tf.Tensor:

    """simple docstring"""







    _snake_case      :					List[Any]					=   hidden_state
    _snake_case      :					Tuple					=   self.conva(lowerCAmelCase			,						training=lowerCAmelCase)
    _snake_case      :					List[Any]					=   self.conva(lowerCAmelCase			,						training=lowerCAmelCase)
    _snake_case      :					str					=   self.shortcut(lowerCAmelCase			,						training=lowerCAmelCase)
    hidden_state += residual
    _snake_case      :					Dict					=   self.activation(lowerCAmelCase)
    return hidden_state







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



   '''simple docstring'''






   def __init__(   self							:     Union[str, Any]			,						lowerCAmelCase							:     int			,						lowerCAmelCase							:     int			,						lowerCAmelCase							:     int = 1			,						lowerCAmelCase							:     str = "relu"			,						lowerCAmelCase							:     int = 4			,						**lowerCAmelCase							:     Dict			,						)			->			None:

    """simple docstring"""







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


   def        UpperCamelCase_					(   self							:     List[Any]			,						lowerCAmelCase							:     tf.Tensor			,						lowerCAmelCase							:     bool = False)			->			tf.Tensor:

    """simple docstring"""







    _snake_case      :					Optional[int]					=   hidden_state
    _snake_case      :					Dict					=   self.conva(lowerCAmelCase			,						training=lowerCAmelCase)
    _snake_case      :					Any					=   self.conva(lowerCAmelCase			,						training=lowerCAmelCase)
    _snake_case      :					Optional[int]					=   self.conva(lowerCAmelCase			,						training=lowerCAmelCase)
    _snake_case      :					str					=   self.shortcut(lowerCAmelCase			,						training=lowerCAmelCase)
    hidden_state += residual
    _snake_case      :					Tuple					=   self.activation(lowerCAmelCase)
    return hidden_state







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



   '''simple docstring'''






   def __init__(   self							:     Optional[Any]			,						lowerCAmelCase							:     ResNetConfig			,						lowerCAmelCase							:     int			,						lowerCAmelCase							:     int			,						lowerCAmelCase							:     int = 2			,						lowerCAmelCase							:     int = 2			,						**lowerCAmelCase							:     Optional[int])			->			None:

    """simple docstring"""







    super().__init__(**lowerCAmelCase)

    _snake_case      :					str					=   TFResNetBottleNeckLayer if config.layer_type == """bottleneck""" else TFResNetBasicLayer

    _snake_case      :					Dict					=   [layer(lowerCAmelCase			,						lowerCAmelCase			,						stride=lowerCAmelCase			,						activation=config.hidden_act			,						name="""layers.0""")]
    layers += [
        layer(lowerCAmelCase			,						lowerCAmelCase			,						activation=config.hidden_act			,						name=F'''layers.{i + 1}''')
        for i in range(depth - 1)
    ]
    _snake_case      :					List[str]					=   layers


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

    """simple docstring"""







    for layer in self.stage_layers:
     _snake_case      :					int					=   layer(lowerCAmelCase			,						training=lowerCAmelCase)
    return hidden_state







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



   '''simple docstring'''






   def __init__(   self							:     List[str]			,						lowerCAmelCase							:     ResNetConfig			,						**lowerCAmelCase							:     List[Any])			->			None:

    """simple docstring"""







    super().__init__(**lowerCAmelCase)
    # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
    _snake_case      :					Tuple					=   [
        TFResNetStage(
            lowerCAmelCase			,						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(lowerCAmelCase			,						lowerCAmelCase			,						lowerCAmelCase			,						depth=lowerCAmelCase			,						name=F'''stages.{i + 1}'''))


   def        UpperCamelCase_					(   self							:     int			,						lowerCAmelCase							:     tf.Tensor			,						lowerCAmelCase							:     bool = False			,						lowerCAmelCase							:     bool = True			,						lowerCAmelCase							:     bool = False			,						)			->			TFBaseModelOutputWithNoAttention:

    """simple docstring"""







    _snake_case      :					List[Any]					=   () if output_hidden_states else None

    for stage_module in self.stages:
     if output_hidden_states:
      _snake_case      :					Optional[Any]					=   hidden_states + (hidden_state,)

     _snake_case      :					Union[str, Any]					=   stage_module(lowerCAmelCase			,						training=lowerCAmelCase)

    if output_hidden_states:
     _snake_case      :					Any					=   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=lowerCAmelCase			,						hidden_states=lowerCAmelCase)







class 						snake_case  (		SCREAMING_SNAKE_CASE_					):



   '''simple docstring'''






   snake_case_   :	Optional[int]       =						ResNetConfig
   snake_case_   :	Any       =						"""resnet"""
   snake_case_   :	Union[str, Any]       =						"""pixel_values"""

   @property
   def        UpperCamelCase_					(   self							:     int)			->			Dict:

    """simple docstring"""







    return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)			,						dtype=tf.floataa)}
a__      =		R"""
    This model is a TensorFlow
    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a
    regular TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
"""


a__      =		R"""
    Args:
        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""







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



   '''simple docstring'''






   snake_case_   :	Optional[int]       =						ResNetConfig

   def __init__(   self							:     Union[str, Any]			,						lowerCAmelCase							:     ResNetConfig			,						**lowerCAmelCase							:     Tuple)			->			None:

    """simple docstring"""







    super().__init__(**lowerCAmelCase)
    _snake_case      :					int					=   config
    _snake_case      :					Tuple					=   TFResNetEmbeddings(lowerCAmelCase			,						name="""embedder""")
    _snake_case      :					str					=   TFResNetEncoder(lowerCAmelCase			,						name="""encoder""")
    _snake_case      :					Any					=   tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowerCAmelCase)


   @unpack_inputs
   def        UpperCamelCase_					(   self							:     str			,						lowerCAmelCase							:     tf.Tensor			,						lowerCAmelCase							:     Optional[bool] = None			,						lowerCAmelCase							:     Optional[bool] = None			,						lowerCAmelCase							:     bool = False			,						)			->			Union[Tuple[tf.Tensor], TFBaseModelOutputWithPoolingAndNoAttention]:

    """simple docstring"""







    _snake_case      :					int					=   (
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    _snake_case      :					Optional[int]					=   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)
    _snake_case      :					str					=   tf.transpose(lowerCAmelCase			,						perm=[0, 2, 3, 1])
    _snake_case      :					Optional[Any]					=   self.embedder(lowerCAmelCase			,						training=lowerCAmelCase)

    _snake_case      :					List[str]					=   self.encoder(
        lowerCAmelCase			,						output_hidden_states=lowerCAmelCase			,						return_dict=lowerCAmelCase			,						training=lowerCAmelCase)

    _snake_case      :					Dict					=   encoder_outputs[0]

    _snake_case      :					Optional[int]					=   self.pooler(lowerCAmelCase)

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

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

    _snake_case      :					Optional[Any]					=   hidden_states if output_hidden_states else None

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







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



   '''simple docstring'''






   def __init__(   self							:     Optional[Any]			,						lowerCAmelCase							:     ResNetConfig			,						**lowerCAmelCase							:     Any)			->			None:

    """simple docstring"""







    super().__init__(lowerCAmelCase			,						**lowerCAmelCase)
    _snake_case      :					Union[str, Any]					=   TFResNetMainLayer(config=lowerCAmelCase			,						name="""resnet""")


   @add_start_docstrings_to_model_forward(lowerCAmelCase)
   @add_code_sample_docstrings(
       checkpoint=_CHECKPOINT_FOR_DOC			,						output_type=lowerCAmelCase			,						config_class=_CONFIG_FOR_DOC			,						modality="""vision"""			,						expected_output=_EXPECTED_OUTPUT_SHAPE			,						)
   @unpack_inputs
   def        UpperCamelCase_					(   self							:     Tuple			,						lowerCAmelCase							:     tf.Tensor			,						lowerCAmelCase							:     Optional[bool] = None			,						lowerCAmelCase							:     Optional[bool] = None			,						lowerCAmelCase							:     bool = False			,						)			->			Union[Tuple[tf.Tensor], TFBaseModelOutputWithPoolingAndNoAttention]:

    """simple docstring"""







    _snake_case      :					List[str]					=   (
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    _snake_case      :					Tuple					=   return_dict if return_dict is not None else self.config.use_return_dict

    _snake_case      :					Union[str, Any]					=   self.resnet(
        pixel_values=lowerCAmelCase			,						output_hidden_states=lowerCAmelCase			,						return_dict=lowerCAmelCase			,						training=lowerCAmelCase			,						)
    return resnet_outputs







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



   '''simple docstring'''






   def __init__(   self							:     Optional[Any]			,						lowerCAmelCase							:     ResNetConfig			,						**lowerCAmelCase							:     Dict)			->			None:

    """simple docstring"""







    super().__init__(lowerCAmelCase			,						**lowerCAmelCase)
    _snake_case      :					Optional[Any]					=   config.num_labels
    _snake_case      :					int					=   TFResNetMainLayer(lowerCAmelCase			,						name="""resnet""")
    # classification head
    _snake_case      :					Any					=   (
        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]			,						lowerCAmelCase							:     tf.Tensor)			->			tf.Tensor:

    """simple docstring"""







    _snake_case      :					List[str]					=   tf.keras.layers.Flatten()(lowerCAmelCase)
    _snake_case      :					List[str]					=   self.classifier_layer(lowerCAmelCase)
    return logits


   @add_start_docstrings_to_model_forward(lowerCAmelCase)
   @add_code_sample_docstrings(
       checkpoint=_IMAGE_CLASS_CHECKPOINT			,						output_type=lowerCAmelCase			,						config_class=_CONFIG_FOR_DOC			,						expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT			,						)
   @unpack_inputs
   def        UpperCamelCase_					(   self							:     List[str]			,						lowerCAmelCase							:     tf.Tensor = None			,						lowerCAmelCase							:     tf.Tensor = None			,						lowerCAmelCase							:     bool = None			,						lowerCAmelCase							:     bool = None			,						lowerCAmelCase							:     bool = False			,						)			->			Union[Tuple[tf.Tensor], TFImageClassifierOutputWithNoAttention]:

    """simple docstring"""







    _snake_case      :					Any					=   return_dict if return_dict is not None else self.config.use_return_dict

    _snake_case      :					Dict					=   self.resnet(
        lowerCAmelCase			,						output_hidden_states=lowerCAmelCase			,						return_dict=lowerCAmelCase			,						training=lowerCAmelCase)

    _snake_case      :					int					=   outputs.pooler_output if return_dict else outputs[1]

    _snake_case      :					List[str]					=   self.classifier(lowerCAmelCase)

    _snake_case      :					Optional[int]					=   None if labels is None else self.hf_compute_loss(lowerCAmelCase			,						lowerCAmelCase)

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

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