





import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import ClassLabel, Features, Image
from .base import TaskTemplate







@dataclass(frozen=lowerCamelCase__  )
class 			__A       (      lowerCamelCase__  ):




      """simple docstring"""







      UpperCAmelCase__						=      field(default="""image-classification""" ,metadata={"""include_in_asdict_even_if_is_default""": True}  )
      UpperCAmelCase__						=      Features({"""image""": Image()}  )
      UpperCAmelCase__						=      Features({"""labels""": ClassLabel}  )
      UpperCAmelCase__						=      "image"
      UpperCAmelCase__						=      "labels"


      def 				__snake_case			(		self      ,							a__):



             """simple docstring"""





             if self.label_column not in features:
                    raise ValueError(F"""Column {self.label_column} is not present in features.""")
             if not isinstance(features[self.label_column]      ,							a__):
                    raise ValueError(F"""Column {self.label_column} is not a ClassLabel.""")
             _lowerCamelCase							:     List[Any]	=		copy.deepcopy(self)
             _lowerCamelCase							:     Optional[int]	=		self.label_schema.copy()
             _lowerCamelCase							:     int	=		features[self.label_column]
             _lowerCamelCase							:     List[str]	=		label_schema
             return task_template



      @property
      def 				__snake_case			(		self):



             """simple docstring"""





             return {
                 self.image_column: "image",
                 self.label_column: "labels",
             }








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, 1088, 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  ):




      """simple docstring"""







      def __init__(		self      ,							a__      ,							a__ = 3      ,							a__ = 1      ,							a__ = 1      ,							a__ = "relu"      ,							**a__      ,							):



             """simple docstring"""





             super().__init__(**a__)
             # The padding and conv has been verified in
             # https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
             _lowerCamelCase							:     List[str]	=		tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2)
             _lowerCamelCase							:     List[str]	=		tf.keras.layers.ConvaD(
                 filters=a__      ,							kernel_size=a__      ,							strides=a__      ,							padding='''VALID'''      ,							groups=a__      ,							use_bias=a__      ,							name='''convolution'''      ,							)
             _lowerCamelCase							:     List[str]	=		tf.keras.layers.BatchNormalization(epsilon=1e-5      ,							momentum=0.9      ,							name='''normalization''')
             _lowerCamelCase							:     Tuple	=		ACTaFN[activation] if activation is not None else tf.identity



      def 				__snake_case			(		self      ,							a__):



             """simple docstring"""





             _lowerCamelCase							:     List[Any]	=		self.convolution(self.padding(a__))
             _lowerCamelCase							:     Any	=		self.normalization(a__)
             _lowerCamelCase							:     Any	=		self.activation(a__)
             return hidden_state







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




      """simple docstring"""







      def __init__(		self      ,							a__      ,							**a__):



             """simple docstring"""





             super().__init__(**a__)
             _lowerCamelCase							:     Optional[int]	=		config.num_channels
             _lowerCamelCase							:     Any	=		TFRegNetConvLayer(
                 out_channels=config.embedding_size      ,							kernel_size=3      ,							stride=2      ,							activation=config.hidden_act      ,							name='''embedder'''      ,							)



      def 				__snake_case			(		self      ,							a__):



             """simple docstring"""





             _lowerCamelCase							:     int	=		shape_list(a__)[1]
             if tf.executing_eagerly() and num_channels != self.num_channels:
                    raise ValueError(
                        '''Make sure that the channel dimension of the pixel values match with the one set in the configuration.''')

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







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




      """simple docstring"""







      def __init__(		self      ,							a__      ,							a__ = 2      ,							**a__):



             """simple docstring"""





             super().__init__(**a__)
             _lowerCamelCase							:     Dict	=		tf.keras.layers.ConvaD(
                 filters=a__      ,							kernel_size=1      ,							strides=a__      ,							use_bias=a__      ,							name='''convolution''')
             _lowerCamelCase							:     List[Any]	=		tf.keras.layers.BatchNormalization(epsilon=1e-5      ,							momentum=0.9      ,							name='''normalization''')



      def 				__snake_case			(		self      ,							a__      ,							a__ = False):



             """simple docstring"""





             return self.normalization(self.convolution(a__)      ,							training=a__)







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




      """simple docstring"""







      def __init__(		self      ,							a__      ,							a__      ,							**a__):



             """simple docstring"""





             super().__init__(**a__)
             _lowerCamelCase							:     List[Any]	=		tf.keras.layers.GlobalAveragePoolingaD(keepdims=a__      ,							name='''pooler''')
             _lowerCamelCase							:     str	=		[
                 tf.keras.layers.ConvaD(filters=a__      ,							kernel_size=1      ,							activation='''relu'''      ,							name='''attention.0'''),
                 tf.keras.layers.ConvaD(filters=a__      ,							kernel_size=1      ,							activation='''sigmoid'''      ,							name='''attention.2'''),
             ]



      def 				__snake_case			(		self      ,							a__):



             """simple docstring"""





             _lowerCamelCase							:     Tuple	=		self.pooler(a__)
             for layer_module in self.attention:
                    _lowerCamelCase							:     Optional[Any]	=		layer_module(a__)
             _lowerCamelCase							:     List[str]	=		hidden_state * pooled
             return hidden_state







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




      """simple docstring"""







      def __init__(		self      ,							a__      ,							a__      ,							a__      ,							a__ = 1      ,							**a__):



             """simple docstring"""





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



      def 				__snake_case			(		self      ,							a__):



             """simple docstring"""





             _lowerCamelCase							:     Optional[Any]	=		hidden_state
             for layer_module in self.layers:
                    _lowerCamelCase							:     Optional[Any]	=		layer_module(a__)
             _lowerCamelCase							:     int	=		self.shortcut(a__)
             hidden_state += residual
             _lowerCamelCase							:     Union[str, Any]	=		self.activation(a__)
             return hidden_state







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




      """simple docstring"""







      def __init__(		self      ,							a__      ,							a__      ,							a__      ,							a__ = 1      ,							**a__):



             """simple docstring"""





             super().__init__(**a__)
             _lowerCamelCase							:     Dict	=		in_channels != out_channels or stride != 1
             _lowerCamelCase							:     str	=		max(1      ,							out_channels // config.groups_width)
             _lowerCamelCase							:     List[str]	=		(
                 TFRegNetShortCut(a__      ,							stride=a__      ,							name='''shortcut''')
                 if should_apply_shortcut
                 else tf.keras.layers.Activation('''linear'''      ,							name='''shortcut''')
             )
             _lowerCamelCase							:     List[Any]	=		[
                 TFRegNetConvLayer(a__      ,							kernel_size=1      ,							activation=config.hidden_act      ,							name='''layer.0'''),
                 TFRegNetConvLayer(
                     a__      ,							stride=a__      ,							groups=a__      ,							activation=config.hidden_act      ,							name='''layer.1'''),
                 TFRegNetSELayer(a__      ,							reduced_channels=int(round(in_channels / 4))      ,							name='''layer.2'''),
                 TFRegNetConvLayer(a__      ,							kernel_size=1      ,							activation=a__      ,							name='''layer.3'''),
             ]
             _lowerCamelCase							:     int	=		ACTaFN[config.hidden_act]



      def 				__snake_case			(		self      ,							a__):



             """simple docstring"""





             _lowerCamelCase							:     List[str]	=		hidden_state
             for layer_module in self.layers:
                    _lowerCamelCase							:     str	=		layer_module(a__)
             _lowerCamelCase							:     str	=		self.shortcut(a__)
             hidden_state += residual
             _lowerCamelCase							:     Optional[Any]	=		self.activation(a__)
             return hidden_state







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




      """simple docstring"""







      def __init__(		self      ,							a__      ,							a__      ,							a__      ,							a__ = 2      ,							a__ = 2      ,							**a__):



             """simple docstring"""





             super().__init__(**a__)

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



      def 				__snake_case			(		self      ,							a__):



             """simple docstring"""





             for layer_module in self.layers:
                    _lowerCamelCase							:     Optional[Any]	=		layer_module(a__)
             return hidden_state







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




      """simple docstring"""







      def __init__(		self      ,							a__      ,							**a__):



             """simple docstring"""





             super().__init__(**a__)
             _lowerCamelCase							:     Optional[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(
                     a__      ,							config.embedding_size      ,							config.hidden_sizes[0]      ,							stride=2 if config.downsample_in_first_stage else 1      ,							depth=config.depths[0]      ,							name='''stages.0'''      ,							))
             _lowerCamelCase							:     List[str]	=		zip(config.hidden_sizes      ,							config.hidden_sizes[1:])
             for i, ((in_channels, out_channels), depth) in enumerate(zip(a__      ,							config.depths[1:])):
                    self.stages.append(TFRegNetStage(a__      ,							a__      ,							a__      ,							depth=a__      ,							name=F"""stages.{i+1}"""))



      def 				__snake_case			(		self      ,							a__      ,							a__ = False      ,							a__ = True):



             """simple docstring"""





             _lowerCamelCase							:     Dict	=		() if output_hidden_states else None

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

                    _lowerCamelCase							:     str	=		stage_module(a__)

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







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




      """simple docstring"""







      UpperCAmelCase__						=      RegNetConfig


      def __init__(		self      ,							a__      ,							**a__):



             """simple docstring"""





             super().__init__(**a__)
             _lowerCamelCase							:     int	=		config
             _lowerCamelCase							:     Tuple	=		TFRegNetEmbeddings(a__      ,							name='''embedder''')
             _lowerCamelCase							:     str	=		TFRegNetEncoder(a__      ,							name='''encoder''')
             _lowerCamelCase							:     int	=		tf.keras.layers.GlobalAveragePoolingaD(keepdims=a__      ,							name='''pooler''')



      @unpack_inputs
      def 				__snake_case			(		self      ,							a__      ,							a__ = None      ,							a__ = None      ,							a__ = False      ,							):



             """simple docstring"""





             _lowerCamelCase							:     Dict	=		(
                 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
             )
             _lowerCamelCase							:     str	=		return_dict if return_dict is not None else self.config.use_return_dict

             _lowerCamelCase							:     Optional[int]	=		self.embedder(a__      ,							training=a__)

             _lowerCamelCase							:     Optional[int]	=		self.encoder(
                 a__      ,							output_hidden_states=a__      ,							return_dict=a__      ,							training=a__)

             _lowerCamelCase							:     Dict	=		encoder_outputs[0]
             _lowerCamelCase							:     Union[str, Any]	=		self.pooler(a__)

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

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

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

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







class 			__A       (      lowerCamelCase__  ):




      """simple docstring"""







      UpperCAmelCase__						=      RegNetConfig
      UpperCAmelCase__						=      """regnet"""
      UpperCAmelCase__						=      """pixel_values"""


      @property
      def 				__snake_case			(		self):



             """simple docstring"""





             return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)      ,							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.""" ,lowerCamelCase__ ,)
class 			__A       (      lowerCamelCase__  ):




      """simple docstring"""







      def __init__(		self      ,							a__      ,							*a__      ,							**a__):



             """simple docstring"""





             super().__init__(a__      ,							*a__      ,							**a__)
             _lowerCamelCase							:     str	=		TFRegNetMainLayer(a__      ,							name='''regnet''')



      @unpack_inputs
      @add_start_docstrings_to_model_forward(a__)
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC      ,							output_type=a__      ,							config_class=_CONFIG_FOR_DOC      ,							modality='''vision'''      ,							expected_output=_EXPECTED_OUTPUT_SHAPE      ,							)
      def 				__snake_case			(		self      ,							a__      ,							a__ = None      ,							a__ = None      ,							a__=False      ,							):



             """simple docstring"""





             _lowerCamelCase							:     str	=		(
                 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
             )
             _lowerCamelCase							:     Any	=		return_dict if return_dict is not None else self.config.use_return_dict

             _lowerCamelCase							:     Tuple	=		self.regnet(
                 pixel_values=a__      ,							output_hidden_states=a__      ,							return_dict=a__      ,							training=a__      ,							)
             if not return_dict:
                    return (outputs[0],) + outputs[1:]

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







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




      """simple docstring"""







      def __init__(		self      ,							a__      ,							*a__      ,							**a__):



             """simple docstring"""





             super().__init__(a__      ,							*a__      ,							**a__)
             _lowerCamelCase							:     Dict	=		config.num_labels
             _lowerCamelCase							:     Union[str, Any]	=		TFRegNetMainLayer(a__      ,							name='''regnet''')
             # classification head
             _lowerCamelCase							:     int	=		[
                 tf.keras.layers.Flatten(),
                 tf.keras.layers.Dense(config.num_labels      ,							name='''classifier.1''') if config.num_labels > 0 else tf.identity,
             ]



      @unpack_inputs
      @add_start_docstrings_to_model_forward(a__)
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT      ,							output_type=a__      ,							config_class=_CONFIG_FOR_DOC      ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,							)
      def 				__snake_case			(		self      ,							a__ = None      ,							a__ = None      ,							a__ = None      ,							a__ = None      ,							a__=False      ,							):



             """simple docstring"""





             _lowerCamelCase							:     Any	=		(
                 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
             )
             _lowerCamelCase							:     str	=		return_dict if return_dict is not None else self.config.use_return_dict

             _lowerCamelCase							:     Union[str, Any]	=		self.regnet(
                 a__      ,							output_hidden_states=a__      ,							return_dict=a__      ,							training=a__)

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

             _lowerCamelCase							:     Union[str, Any]	=		self.classifier[0](a__)
             _lowerCamelCase							:     str	=		self.classifier[1](a__)

             _lowerCamelCase							:     List[Any]	=		None if labels is None else self.hf_compute_loss(labels=a__      ,							logits=a__)

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

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

