






"""simple docstring"""





import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing the experiment tracking capability,
# and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################

__lowerCamelCase							=  16
__lowerCamelCase							=  32



def   a    (     __UpperCAmelCase      :    Accelerator      ,   __UpperCAmelCase      :    int = 1_6  )    ->							int:
						__magic_name__:	Optional[int]							=       AutoTokenizer.from_pretrained("""bert-base-cased"""  )
						__magic_name__:	Any							=       load_dataset("""glue"""      ,   """mrpc"""  )

						def tokenize_function(__UpperCAmelCase      :    Optional[Any]  ):
												# max_length=None => use the model max length (it's actually the default)
												__magic_name__:	Dict							=       tokenizer(examples["""sentence1"""]      ,   examples["""sentence2"""]      ,   truncation=__UpperCAmelCase      ,   max_length=__UpperCAmelCase  )
												return outputs

						# Apply the method we just defined to all the examples in all the splits of the dataset
						# starting with the main process first:
						with accelerator.main_process_first():
												__magic_name__:	int							=       datasets.map(
												    __UpperCAmelCase      ,   batched=__UpperCAmelCase      ,   remove_columns=["""idx""", """sentence1""", """sentence2"""]      ,   )

						# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
						# transformers library
						__magic_name__:	Tuple							=       tokenized_datasets.rename_column("""label"""      ,   """labels"""  )

						def collate_fn(__UpperCAmelCase      :    Union[str, Any]  ):
												# On TPU it's best to pad everything to the same length or training will be very slow.
												__magic_name__:	str							=       1_2_8 if accelerator.distributed_type == DistributedType.TPU else None
												# When using mixed precision we want round multiples of 8/16
												if accelerator.mixed_precision == "fp8":
																		__magic_name__:	Union[str, Any]							=       1_6
												elif accelerator.mixed_precision != "no":
																		__magic_name__:	List[str]							=       8
												else:
																		__magic_name__:	int							=       None

												return tokenizer.pad(
												    __UpperCAmelCase      ,   padding="""longest"""      ,   max_length=__UpperCAmelCase      ,   pad_to_multiple_of=__UpperCAmelCase      ,   return_tensors="""pt"""      ,   )

						# Instantiate dataloaders.
						__magic_name__:	List[Any]							=       DataLoader(
						    tokenized_datasets["""train"""]      ,   shuffle=__UpperCAmelCase      ,   collate_fn=__UpperCAmelCase      ,   batch_size=__UpperCAmelCase  )
						__magic_name__:	str							=       DataLoader(
						    tokenized_datasets["""validation"""]      ,   shuffle=__UpperCAmelCase      ,   collate_fn=__UpperCAmelCase      ,   batch_size=__UpperCAmelCase  )

						return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('TESTING_MOCKED_DATALOADERS', None) == "1":
							from accelerate.test_utils.training import mocked_dataloaders

							__lowerCamelCase							=  mocked_dataloaders  # noqa: F811



def   a    (     __UpperCAmelCase      :    Tuple      ,   __UpperCAmelCase      :    Union[str, Any]  )    ->							Optional[int]:
						# For testing only
						if os.environ.get("""TESTING_MOCKED_DATALOADERS"""      ,   __UpperCAmelCase  ) == "1":
												__magic_name__:	int							=       2
						# Initialize Accelerator

						# New Code #
						# We pass in "all" to `log_with` to grab all available trackers in the environment
						# Note: If using a custom `Tracker` class, should be passed in here such as:
						# >>> log_with = ["all", MyCustomTrackerClassInstance()]
						if args.with_tracking:
												__magic_name__:	Union[str, Any]							=       Accelerator(
												    cpu=args.cpu      ,   mixed_precision=args.mixed_precision      ,   log_with="""all"""      ,   project_dir=args.project_dir  )
						else:
												__magic_name__:	Tuple							=       Accelerator(cpu=args.cpu      ,   mixed_precision=args.mixed_precision  )
						# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
						__magic_name__:	str							=       config["""lr"""]
						__magic_name__:	Union[str, Any]							=       int(config["""num_epochs"""]  )
						__magic_name__:	List[str]							=       int(config["""seed"""]  )
						__magic_name__:	Union[str, Any]							=       int(config["""batch_size"""]  )
						set_seed(__UpperCAmelCase  )

						__magic_name__,			__magic_name__:	Tuple							=       get_dataloaders(__UpperCAmelCase      ,   __UpperCAmelCase  )
						__magic_name__:	str							=       evaluate.load("""glue"""      ,   """mrpc"""  )

						# If the batch size is too big we use gradient accumulation
						__magic_name__:	Optional[int]							=       1
						if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
												__magic_name__:	Optional[Any]							=       batch_size // MAX_GPU_BATCH_SIZE
												__magic_name__:	Dict							=       MAX_GPU_BATCH_SIZE

						# Instantiate the model (we build the model here so that the seed also control new weights initialization)
						__magic_name__:	Any							=       AutoModelForSequenceClassification.from_pretrained("""bert-base-cased"""      ,   return_dict=__UpperCAmelCase  )

						# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
						# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
						# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
						__magic_name__:	Dict							=       model.to(accelerator.device  )

						# Instantiate optimizer
						__magic_name__:	List[str]							=       AdamW(params=model.parameters()      ,   lr=__UpperCAmelCase  )

						# Instantiate scheduler
						__magic_name__:	Optional[Any]							=       get_linear_schedule_with_warmup(
						    optimizer=__UpperCAmelCase      ,   num_warmup_steps=1_0_0      ,   num_training_steps=(len(__UpperCAmelCase  ) * num_epochs) // gradient_accumulation_steps      ,   )

						# Prepare everything
						# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
						# prepare method.
						__magic_name__,			__magic_name__,			__magic_name__,			__magic_name__,			__magic_name__:	List[str]							=       accelerator.prepare(
						    __UpperCAmelCase      ,   __UpperCAmelCase      ,   __UpperCAmelCase      ,   __UpperCAmelCase      ,   __UpperCAmelCase  )

						# New Code #
						# We need to initialize the trackers we use. Overall configurations can also be stored
						if args.with_tracking:
												__magic_name__:	List[str]							=       os.path.split(__UpperCAmelCase  )[-1].split("""."""  )[0]
												accelerator.init_trackers(__UpperCAmelCase      ,   __UpperCAmelCase  )

						# Now we train the model
						for epoch in range(__UpperCAmelCase  ):
												model.train()
												# New Code #
												# For our tracking example, we will log the total loss of each epoch
												if args.with_tracking:
																		__magic_name__:	int							=       0
												for step, batch in enumerate(__UpperCAmelCase  ):
																		# We could avoid this line since we set the accelerator with `device_placement=True`.
																		batch.to(accelerator.device  )
																		__magic_name__:	Any							=       model(**__UpperCAmelCase  )
																		__magic_name__:	int							=       outputs.loss
																		# New Code #
																		if args.with_tracking:
																								total_loss += loss.detach().float()
																		__magic_name__:	Tuple							=       loss / gradient_accumulation_steps
																		accelerator.backward(__UpperCAmelCase  )
																		if step % gradient_accumulation_steps == 0:
																								optimizer.step()
																								lr_scheduler.step()
																								optimizer.zero_grad()

												model.eval()
												for step, batch in enumerate(__UpperCAmelCase  ):
																		# We could avoid this line since we set the accelerator with `device_placement=True` (the default).
																		batch.to(accelerator.device  )
																		with torch.no_grad():
																								__magic_name__:	int							=       model(**__UpperCAmelCase  )
																		__magic_name__:	int							=       outputs.logits.argmax(dim=-1  )
																		__magic_name__,			__magic_name__:	List[str]							=       accelerator.gather_for_metrics((predictions, batch["""labels"""])  )
																		metric.add_batch(
																		    predictions=__UpperCAmelCase      ,   references=__UpperCAmelCase      ,   )

												__magic_name__:	Union[str, Any]							=       metric.compute()
												# Use accelerator.print to print only on the main process.
												accelerator.print(f'epoch {epoch}:'      ,   __UpperCAmelCase  )

												# New Code #
												# To actually log, we call `Accelerator.log`
												# The values passed can be of `str`, `int`, `float` or `dict` of `str` to `float`/`int`
												if args.with_tracking:
																		accelerator.log(
																		    {
																		        """accuracy""": eval_metric["""accuracy"""],
																		        """f1""": eval_metric["""f1"""],
																		        """train_loss""": total_loss.item() / len(__UpperCAmelCase  ),
																		        """epoch""": epoch,
																		    }      ,   step=__UpperCAmelCase      ,   )

    # New Code #
    # When a run is finished, you should call `accelerator.end_training()`
    # to close all of the open trackers
						if args.with_tracking:
												accelerator.end_training()



def   a    (     )    ->							str:
						__magic_name__:	Optional[int]							=       argparse.ArgumentParser(description="""Simple example of training script."""  )
						parser.add_argument(
						    """--mixed_precision"""      ,   type=__UpperCAmelCase      ,   default=__UpperCAmelCase      ,   choices=["""no""", """fp16""", """bf16""", """fp8"""]      ,   help="""Whether to use mixed precision. Choose"""
						    """between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."""
						    """and an Nvidia Ampere GPU."""      ,   )
						parser.add_argument("""--cpu"""      ,   action="""store_true"""      ,   help="""If passed, will train on the CPU."""  )
						parser.add_argument(
						    """--with_tracking"""      ,   action="""store_true"""      ,   help="""Whether to load in all available experiment trackers from the environment and use them for logging."""      ,   )
						parser.add_argument(
						    """--project_dir"""      ,   type=__UpperCAmelCase      ,   default="""logs"""      ,   help="""Location on where to store experiment tracking logs` and relevent project information"""      ,   )
						__magic_name__:	Union[str, Any]							=       parser.parse_args()
						__magic_name__:	int							=       {"""lr""": 2E-5, """num_epochs""": 3, """seed""": 4_2, """batch_size""": 1_6}
						training_function(__UpperCAmelCase      ,   __UpperCAmelCase  )


if __name__ == "__main__":
							main()











"""simple docstring"""





from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


__lowerCamelCase							=  logging.get_logger(__name__)

# General docstring
__lowerCamelCase							=  'RegNetConfig'

# Base docstring
__lowerCamelCase							=  'facebook/regnet-y-040'
__lowerCamelCase							=  [1, 10_88, 7, 7]

# Image classification docstring
__lowerCamelCase							=  'facebook/regnet-y-040'
__lowerCamelCase							=  'tabby, tabby cat'

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



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

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






					def   lowerCamelCase__   (			self  :      Optional[int]     ,		__snake_case  :      str					)			->							Dict:
											__magic_name__:	Optional[Any]							=       self.convolution(self.padding(__snake_case					)					)
											__magic_name__:	Union[str, Any]							=       self.normalization(__snake_case					)
											__magic_name__:	Tuple							=       self.activation(__snake_case					)
											return hidden_state



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

					def __init__(			self  :      Union[str, Any]     ,		__snake_case  :      RegNetConfig     ,		**__snake_case  :      Dict					)			->							Optional[int]:
											super().__init__(**__snake_case					)
											__magic_name__:	Tuple							=       config.num_channels
											__magic_name__:	Optional[int]							=       TFRegNetConvLayer(
											    out_channels=config.embedding_size     ,		kernel_size=3     ,		stride=2     ,		activation=config.hidden_act     ,		name="""embedder"""     ,		)






					def   lowerCamelCase__   (			self  :      List[str]     ,		__snake_case  :      Dict					)			->							int:
											__magic_name__:	Union[str, Any]							=       shape_list(__snake_case					)[1]
											if tf.executing_eagerly() and num_channels != self.num_channels:
																	raise ValueError(
																	    """Make sure that the channel dimension of the pixel values match with the one set in the configuration."""					)

											# When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
											# So change the input format from `NCHW` to `NHWC`.
											# shape = (batch_size, in_height, in_width, in_channels=num_channels)
											__magic_name__:	Any							=       tf.transpose(__snake_case     ,		perm=(0, 2, 3, 1)					)
											__magic_name__:	Dict							=       self.embedder(__snake_case					)
											return hidden_state



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

					def __init__(			self  :      Union[str, Any]     ,		__snake_case  :      int     ,		__snake_case  :      int = 2     ,		**__snake_case  :      Any					)			->							Dict:
											super().__init__(**__snake_case					)
											__magic_name__:	Union[str, Any]							=       tf.keras.layers.ConvaD(
											    filters=__snake_case     ,		kernel_size=1     ,		strides=__snake_case     ,		use_bias=__snake_case     ,		name="""convolution"""					)
											__magic_name__:	Dict							=       tf.keras.layers.BatchNormalization(epsilon=1E-5     ,		momentum=0.9     ,		name="""normalization"""					)






					def   lowerCamelCase__   (			self  :      Optional[Any]     ,		__snake_case  :      tf.Tensor     ,		__snake_case  :      bool = False					)			->							tf.Tensor:
											return self.normalization(self.convolution(__snake_case					)     ,		training=__snake_case					)



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

					def __init__(			self  :      int     ,		__snake_case  :      int     ,		__snake_case  :      int     ,		**__snake_case  :      str					)			->							str:
											super().__init__(**__snake_case					)
											__magic_name__:	Optional[Any]							=       tf.keras.layers.GlobalAveragePoolingaD(keepdims=__snake_case     ,		name="""pooler"""					)
											__magic_name__:	Optional[Any]							=       [
											    tf.keras.layers.ConvaD(filters=__snake_case     ,		kernel_size=1     ,		activation="""relu"""     ,		name="""attention.0"""					),
											    tf.keras.layers.ConvaD(filters=__snake_case     ,		kernel_size=1     ,		activation="""sigmoid"""     ,		name="""attention.2"""					),
											]






					def   lowerCamelCase__   (			self  :      Dict     ,		__snake_case  :      List[str]					)			->							List[Any]:
											# [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
											__magic_name__:	List[str]							=       self.pooler(__snake_case					)
											for layer_module in self.attention:
																	__magic_name__:	List[str]							=       layer_module(__snake_case					)
											__magic_name__:	Optional[Any]							=       hidden_state * pooled
											return hidden_state



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

					def __init__(			self  :      Union[str, Any]     ,		__snake_case  :      RegNetConfig     ,		__snake_case  :      int     ,		__snake_case  :      int     ,		__snake_case  :      int = 1     ,		**__snake_case  :      Optional[int]					)			->							Optional[int]:
											super().__init__(**__snake_case					)
											__magic_name__:	List[str]							=       in_channels != out_channels or stride != 1
											__magic_name__:	Union[str, Any]							=       max(1     ,		out_channels // config.groups_width					)
											__magic_name__:	Optional[Any]							=       (
											    TFRegNetShortCut(__snake_case     ,		stride=__snake_case     ,		name="""shortcut"""					)
											    if should_apply_shortcut
											    else tf.keras.layers.Activation("""linear"""     ,		name="""shortcut"""					)
											)
											# `self.layers` instead of `self.layer` because that is a reserved argument.
											__magic_name__:	List[str]							=       [
											    TFRegNetConvLayer(__snake_case     ,		kernel_size=1     ,		activation=config.hidden_act     ,		name="""layer.0"""					),
											    TFRegNetConvLayer(
											        __snake_case     ,		stride=__snake_case     ,		groups=__snake_case     ,		activation=config.hidden_act     ,		name="""layer.1"""					),
											    TFRegNetConvLayer(__snake_case     ,		kernel_size=1     ,		activation=__snake_case     ,		name="""layer.2"""					),
											]
											__magic_name__:	Any							=       ACTaFN[config.hidden_act]






					def   lowerCamelCase__   (			self  :      Optional[int]     ,		__snake_case  :      Any					)			->							Union[str, Any]:
											__magic_name__:	Any							=       hidden_state
											for layer_module in self.layers:
																	__magic_name__:	Optional[int]							=       layer_module(__snake_case					)
											__magic_name__:	str							=       self.shortcut(__snake_case					)
											hidden_state += residual
											__magic_name__:	int							=       self.activation(__snake_case					)
											return hidden_state



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

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






					def   lowerCamelCase__   (			self  :      List[str]     ,		__snake_case  :      int					)			->							Dict:
											__magic_name__:	int							=       hidden_state
											for layer_module in self.layers:
																	__magic_name__:	Optional[Any]							=       layer_module(__snake_case					)
											__magic_name__:	Union[str, Any]							=       self.shortcut(__snake_case					)
											hidden_state += residual
											__magic_name__:	Any							=       self.activation(__snake_case					)
											return hidden_state



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

					def __init__(			self  :      int     ,		__snake_case  :      RegNetConfig     ,		__snake_case  :      int     ,		__snake_case  :      int     ,		__snake_case  :      int = 2     ,		__snake_case  :      int = 2     ,		**__snake_case  :      List[Any]					)			->							Optional[int]:
											super().__init__(**__snake_case					)

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






					def   lowerCamelCase__   (			self  :      int     ,		__snake_case  :      Union[str, Any]					)			->							Tuple:
											for layer_module in self.layers:
																	__magic_name__:	Dict							=       layer_module(__snake_case					)
											return hidden_state



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

					def __init__(			self  :      Union[str, Any]     ,		__snake_case  :      RegNetConfig     ,		**__snake_case  :      Optional[Any]					)			->							Dict:
											super().__init__(**__snake_case					)
											__magic_name__:	List[Any]							=       []
											# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
											self.stages.append(
											    TFRegNetStage(
											        __snake_case     ,		config.embedding_size     ,		config.hidden_sizes[0]     ,		stride=2 if config.downsample_in_first_stage else 1     ,		depth=config.depths[0]     ,		name="""stages.0"""     ,		)					)
											__magic_name__:	Dict							=       zip(config.hidden_sizes     ,		config.hidden_sizes[1:]					)
											for i, ((in_channels, out_channels), depth) in enumerate(zip(__snake_case     ,		config.depths[1:]					)					):
																	self.stages.append(TFRegNetStage(__snake_case     ,		__snake_case     ,		__snake_case     ,		depth=__snake_case     ,		name=F'stages.{i+1}'					)					)






					def   lowerCamelCase__   (			self  :      int     ,		__snake_case  :      tf.Tensor     ,		__snake_case  :      bool = False     ,		__snake_case  :      bool = True					)			->							TFBaseModelOutputWithNoAttention:
											__magic_name__:	int							=       () if output_hidden_states else None

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

																	__magic_name__:	Optional[Any]							=       stage_module(__snake_case					)

											if output_hidden_states:
																	__magic_name__:	int							=       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=__snake_case     ,		hidden_states=__snake_case					)



@keras_serializable
class       __A		(   tf.keras.layers.Layer			):
					UpperCAmelCase__								=				RegNetConfig

					def __init__(			self  :      Optional[int]     ,		__snake_case  :      Any     ,		**__snake_case  :      List[str]					)			->							int:
											super().__init__(**__snake_case					)
											__magic_name__:	Union[str, Any]							=       config
											__magic_name__:	Optional[int]							=       TFRegNetEmbeddings(__snake_case     ,		name="""embedder"""					)
											__magic_name__:	int							=       TFRegNetEncoder(__snake_case     ,		name="""encoder"""					)
											__magic_name__:	int							=       tf.keras.layers.GlobalAveragePoolingaD(keepdims=__snake_case     ,		name="""pooler"""					)






					@unpack_inputs
					def   lowerCamelCase__   (			self  :      Optional[Any]     ,		__snake_case  :      tf.Tensor     ,		__snake_case  :      Optional[bool] = None     ,		__snake_case  :      Optional[bool] = None     ,		__snake_case  :      bool = False     ,		)			->							TFBaseModelOutputWithPoolingAndNoAttention:
											__magic_name__:	Optional[int]							=       (
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											__magic_name__:	int							=       return_dict if return_dict is not None else self.config.use_return_dict

											__magic_name__:	List[str]							=       self.embedder(__snake_case     ,		training=__snake_case					)

											__magic_name__:	Optional[Any]							=       self.encoder(
											    __snake_case     ,		output_hidden_states=__snake_case     ,		return_dict=__snake_case     ,		training=__snake_case					)

											__magic_name__:	str							=       encoder_outputs[0]
											__magic_name__:	List[Any]							=       self.pooler(__snake_case					)

											# Change to NCHW output format have uniformity in the modules
											__magic_name__:	int							=       tf.transpose(__snake_case     ,		perm=(0, 3, 1, 2)					)
											__magic_name__:	List[str]							=       tf.transpose(__snake_case     ,		perm=(0, 3, 1, 2)					)

											# Change the other hidden state outputs to NCHW as well
											if output_hidden_states:
																	__magic_name__:	List[str]							=       tuple([tf.transpose(__snake_case     ,		perm=(0, 3, 1, 2)					) for h in encoder_outputs[1]]					)

											if not return_dict:
																	return (last_hidden_state, pooled_output) + encoder_outputs[1:]

											return TFBaseModelOutputWithPoolingAndNoAttention(
											    last_hidden_state=__snake_case     ,		pooler_output=__snake_case     ,		hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states     ,		)



class       __A		(   SCREAMING_SNAKE_CASE_			):
					UpperCAmelCase__								=				RegNetConfig
					UpperCAmelCase__								=				"regnet"
					UpperCAmelCase__								=				"pixel_values"

					@property
					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							Union[str, Any]:
											return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_2_4, 2_2_4)     ,		dtype=tf.floataa					)}


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

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



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

					def __init__(			self  :      Optional[Any]     ,		__snake_case  :      RegNetConfig     ,		*__snake_case  :      List[Any]     ,		**__snake_case  :      Tuple					)			->							Tuple:
											super().__init__(__snake_case     ,		*__snake_case     ,		**__snake_case					)
											__magic_name__:	List[str]							=       TFRegNetMainLayer(__snake_case     ,		name="""regnet"""					)






					@unpack_inputs
					@add_start_docstrings_to_model_forward(__snake_case					)
					@add_code_sample_docstrings(
					    checkpoint=_CHECKPOINT_FOR_DOC     ,		output_type=__snake_case     ,		config_class=_CONFIG_FOR_DOC     ,		modality="""vision"""     ,		expected_output=_EXPECTED_OUTPUT_SHAPE     ,		)
					def   lowerCamelCase__   (			self  :      Dict     ,		__snake_case  :      tf.Tensor     ,		__snake_case  :      Optional[bool] = None     ,		__snake_case  :      Optional[bool] = None     ,		__snake_case  :      int=False     ,		)			->							Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:
											__magic_name__:	Optional[int]							=       (
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											__magic_name__:	Union[str, Any]							=       return_dict if return_dict is not None else self.config.use_return_dict

											__magic_name__:	List[str]							=       self.regnet(
											    pixel_values=__snake_case     ,		output_hidden_states=__snake_case     ,		return_dict=__snake_case     ,		training=__snake_case     ,		)
											if not return_dict:
																	return (outputs[0],) + outputs[1:]

											return TFBaseModelOutputWithPoolingAndNoAttention(
											    last_hidden_state=outputs.last_hidden_state     ,		pooler_output=outputs.pooler_output     ,		hidden_states=outputs.hidden_states     ,		)



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

					def __init__(			self  :      int     ,		__snake_case  :      RegNetConfig     ,		*__snake_case  :      Any     ,		**__snake_case  :      Any					)			->							Optional[Any]:
											super().__init__(__snake_case     ,		*__snake_case     ,		**__snake_case					)
											__magic_name__:	Union[str, Any]							=       config.num_labels
											__magic_name__:	Tuple							=       TFRegNetMainLayer(__snake_case     ,		name="""regnet"""					)
											# classification head
											__magic_name__:	List[Any]							=       [
											    tf.keras.layers.Flatten(),
											    tf.keras.layers.Dense(config.num_labels     ,		name="""classifier.1"""					) if config.num_labels > 0 else tf.identity,
											]






					@unpack_inputs
					@add_start_docstrings_to_model_forward(__snake_case					)
					@add_code_sample_docstrings(
					    checkpoint=_IMAGE_CLASS_CHECKPOINT     ,		output_type=__snake_case     ,		config_class=_CONFIG_FOR_DOC     ,		expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT     ,		)
					def   lowerCamelCase__   (			self  :      List[str]     ,		__snake_case  :      tf.Tensor = None     ,		__snake_case  :      tf.Tensor = None     ,		__snake_case  :      bool = None     ,		__snake_case  :      bool = None     ,		__snake_case  :      Dict=False     ,		)			->							Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
											__magic_name__:	Optional[Any]							=       (
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											__magic_name__:	Optional[Any]							=       return_dict if return_dict is not None else self.config.use_return_dict

											__magic_name__:	Any							=       self.regnet(
											    __snake_case     ,		output_hidden_states=__snake_case     ,		return_dict=__snake_case     ,		training=__snake_case					)

											__magic_name__:	Optional[Any]							=       outputs.pooler_output if return_dict else outputs[1]

											__magic_name__:	Optional[int]							=       self.classifier[0](__snake_case					)
											__magic_name__:	List[Any]							=       self.classifier[1](__snake_case					)

											__magic_name__:	Optional[int]							=       None if labels is None else self.hf_compute_loss(labels=__snake_case     ,		logits=__snake_case					)

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

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



