





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


snake_case_           =      logging.get_logger(__name__)

# General docstring
snake_case_           =      "RegNetConfig"

# Base docstring
snake_case_           =      "facebook/regnet-y-040"
snake_case_           =      [1, 1088, 7, 7]

# Image classification docstring
snake_case_           =      "facebook/regnet-y-040"
snake_case_           =      "tabby, tabby cat"

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


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


      def __init__(self	:				Union[str, Any]    ,      a__	:				Dict    ,      a__	:				List[str] = 3    ,      a__	:				List[str] = 1    ,      a__	:				Union[str, Any] = 1    ,      a__	:				Dict = "relu"    ,      **a__	:				int    ,      ):





         """simple docstring"""



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


      def        a			(self	:				int    ,      a__	:				int				):





         """simple docstring"""



         __snake_case								=       self.convolution(self.padding(A__				)				)
         __snake_case								=       self.normalization(A__				)
         __snake_case								=       self.activation(A__				)
         return hidden_state






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


      def __init__(self	:				Dict    ,      a__	:				List[Any]    ,      **a__	:				str				):





         """simple docstring"""



         super().__init__(**A__				)
         __snake_case								=       config.num_channels
         __snake_case								=       TFRegNetConvLayer(
             out_channels=config.embedding_size    ,      kernel_size=3    ,      stride=2    ,      activation=config.hidden_act    ,      name='''embedder'''    ,      )


      def        a			(self	:				str    ,      a__	:				List[str]				):





         """simple docstring"""



         __snake_case								=       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)
         __snake_case								=       tf.transpose(A__    ,      perm=(0, 2, 3, 1)				)
         __snake_case								=       self.embedder(A__				)
         return hidden_state






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


      def __init__(self	:				Union[str, Any]    ,      a__	:				Dict    ,      a__	:				Optional[Any] = 2    ,      **a__	:				Any				):





         """simple docstring"""



         super().__init__(**A__				)
         __snake_case								=       tf.keras.layers.ConvaD(
             filters=A__    ,      kernel_size=1    ,      strides=A__    ,      use_bias=A__    ,      name='''convolution'''				)
         __snake_case								=       tf.keras.layers.BatchNormalization(epsilon=1E-5    ,      momentum=0.9    ,      name='''normalization'''				)


      def        a			(self	:				Tuple    ,      a__	:				List[str]    ,      a__	:				Dict = False				):





         """simple docstring"""



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






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


      def __init__(self	:				Union[str, Any]    ,      a__	:				Any    ,      a__	:				Any    ,      **a__	:				Optional[Any]				):





         """simple docstring"""



         super().__init__(**A__				)
         __snake_case								=       tf.keras.layers.GlobalAveragePoolingaD(keepdims=A__    ,      name='''pooler'''				)
         __snake_case								=       [
             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        a			(self	:				Dict    ,      a__	:				Dict				):





         """simple docstring"""



         __snake_case								=       self.pooler(A__				)
         for layer_module in self.attention:
            __snake_case								=       layer_module(A__				)
         __snake_case								=       hidden_state * pooled
         return hidden_state






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


      def __init__(self	:				Tuple    ,      a__	:				Optional[int]    ,      a__	:				List[str]    ,      a__	:				List[Any]    ,      a__	:				Optional[Any] = 1    ,      **a__	:				Dict				):





         """simple docstring"""



         super().__init__(**A__				)
         __snake_case								=       in_channels != out_channels or stride != 1
         __snake_case								=       max(1    ,      out_channels // config.groups_width				)
         __snake_case								=       (
             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.
         __snake_case								=       [
             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'''				),
         ]
         __snake_case								=       ACTaFN[config.hidden_act]


      def        a			(self	:				Dict    ,      a__	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       hidden_state
         for layer_module in self.layers:
            __snake_case								=       layer_module(A__				)
         __snake_case								=       self.shortcut(A__				)
         hidden_state += residual
         __snake_case								=       self.activation(A__				)
         return hidden_state






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


      def __init__(self	:				List[Any]    ,      a__	:				Optional[int]    ,      a__	:				Dict    ,      a__	:				int    ,      a__	:				Tuple = 1    ,      **a__	:				Union[str, Any]				):





         """simple docstring"""



         super().__init__(**A__				)
         __snake_case								=       in_channels != out_channels or stride != 1
         __snake_case								=       max(1    ,      out_channels // config.groups_width				)
         __snake_case								=       (
             TFRegNetShortCut(A__    ,      stride=A__    ,      name='''shortcut'''				)
             if should_apply_shortcut
             else tf.keras.layers.Activation('''linear'''    ,      name='''shortcut'''				)
         )
         __snake_case								=       [
             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'''				),
         ]
         __snake_case								=       ACTaFN[config.hidden_act]


      def        a			(self	:				int    ,      a__	:				Dict				):





         """simple docstring"""



         __snake_case								=       hidden_state
         for layer_module in self.layers:
            __snake_case								=       layer_module(A__				)
         __snake_case								=       self.shortcut(A__				)
         hidden_state += residual
         __snake_case								=       self.activation(A__				)
         return hidden_state






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


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





         """simple docstring"""



         super().__init__(**A__				)

         __snake_case								=       TFRegNetXLayer if config.layer_type == '''x''' else TFRegNetYLayer
         __snake_case								=       [
             # 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        a			(self	:				Tuple    ,      a__	:				int				):





         """simple docstring"""



         for layer_module in self.layers:
            __snake_case								=       layer_module(A__				)
         return hidden_state






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


      def __init__(self	:				Tuple    ,      a__	:				List[str]    ,      **a__	:				int				):





         """simple docstring"""



         super().__init__(**A__				)
         __snake_case								=       []
         # 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'''    ,      )				)
         __snake_case								=       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        a			(self	:				Any    ,      a__	:				Any    ,      a__	:				str = False    ,      a__	:				Any = True				):





         """simple docstring"""



         __snake_case								=       () if output_hidden_states else None

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

            __snake_case								=       stage_module(A__				)

         if output_hidden_states:
            __snake_case								=       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 		SCREAMING_SNAKE_CASE__    (	tf.keras.layers.Layer  ):
      A_						:    int							=   RegNetConfig


      def __init__(self	:				Tuple    ,      a__	:				List[str]    ,      **a__	:				Optional[Any]				):





         """simple docstring"""



         super().__init__(**A__				)
         __snake_case								=       config
         __snake_case								=       TFRegNetEmbeddings(A__    ,      name='''embedder'''				)
         __snake_case								=       TFRegNetEncoder(A__    ,      name='''encoder'''				)
         __snake_case								=       tf.keras.layers.GlobalAveragePoolingaD(keepdims=A__    ,      name='''pooler'''				)


      @unpack_inputs
      def        a			(self	:				Optional[int]    ,      a__	:				Union[str, Any]    ,      a__	:				Optional[int] = None    ,      a__	:				List[Any] = None    ,      a__	:				Union[str, Any] = False    ,      ):





         """simple docstring"""



         __snake_case								=       (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         __snake_case								=       return_dict if return_dict is not None else self.config.use_return_dict

         __snake_case								=       self.embedder(A__    ,      training=A__				)

         __snake_case								=       self.encoder(
             A__    ,      output_hidden_states=A__    ,      return_dict=A__    ,      training=A__				)

         __snake_case								=       encoder_outputs[0]
         __snake_case								=       self.pooler(A__				)

         # Change to NCHW output format have uniformity in the modules
         __snake_case								=       tf.transpose(A__    ,      perm=(0, 3, 1, 2)				)
         __snake_case								=       tf.transpose(A__    ,      perm=(0, 3, 1, 2)				)

         # Change the other hidden state outputs to NCHW as well
         if output_hidden_states:
            __snake_case								=       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 		SCREAMING_SNAKE_CASE__    (	SCREAMING_SNAKE_CASE__  ):
      A_						:    List[str]							=   RegNetConfig
      A_						:    Any							=   '''regnet'''
      A_						:    str							=   '''pixel_values'''


      @property
      def        a			(self	:				List[str]				):





         """simple docstring"""



         return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)    ,      dtype=tf.floataa				)}






snake_case_           =      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"

snake_case_           =      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 		SCREAMING_SNAKE_CASE__    (	SCREAMING_SNAKE_CASE__  ):


      def __init__(self	:				Dict    ,      a__	:				Optional[Any]    ,      *a__	:				Any    ,      **a__	:				List[str]				):





         """simple docstring"""



         super().__init__(A__    ,      *A__    ,      **A__				)
         __snake_case								=       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        a			(self	:				int    ,      a__	:				List[str]    ,      a__	:				Union[str, Any] = None    ,      a__	:				Union[str, Any] = None    ,      a__	:				List[str]=False    ,      ):





         """simple docstring"""



         __snake_case								=       (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         __snake_case								=       return_dict if return_dict is not None else self.config.use_return_dict

         __snake_case								=       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(
    '\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 		SCREAMING_SNAKE_CASE__    (	SCREAMING_SNAKE_CASE__						,				SCREAMING_SNAKE_CASE__  ):


      def __init__(self	:				Tuple    ,      a__	:				Union[str, Any]    ,      *a__	:				Any    ,      **a__	:				List[Any]				):





         """simple docstring"""



         super().__init__(A__    ,      *A__    ,      **A__				)
         __snake_case								=       config.num_labels
         __snake_case								=       TFRegNetMainLayer(A__    ,      name='''regnet'''				)
         # classification head
         __snake_case								=       [
             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        a			(self	:				List[str]    ,      a__	:				int = None    ,      a__	:				Union[str, Any] = None    ,      a__	:				Optional[Any] = None    ,      a__	:				Union[str, Any] = None    ,      a__	:				Dict=False    ,      ):





         """simple docstring"""



         __snake_case								=       (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         __snake_case								=       return_dict if return_dict is not None else self.config.use_return_dict

         __snake_case								=       self.regnet(
             A__    ,      output_hidden_states=A__    ,      return_dict=A__    ,      training=A__				)

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

         __snake_case								=       self.classifier[0](A__				)
         __snake_case								=       self.classifier[1](A__				)

         __snake_case								=       None if labels is None else self.hf_compute_loss(labels=A__    ,      logits=A__				)

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

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








# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
    import torch

class 		SCREAMING_SNAKE_CASE__    (	TensorFormatter[Mapping, 'torch.Tensor', Mapping]  ):


      def __init__(self	:				Optional[int]    ,      a__	:				Tuple=None    ,      **a__	:				Optional[int]				):





         """simple docstring"""



         super().__init__(features=a__				)
         __snake_case								=       torch_tensor_kwargs
         import torch  # noqa import torch at initialization


      def        a			(self	:				Union[str, Any]    ,      a__	:				Union[str, Any]				):





         """simple docstring"""



         import torch

         if isinstance(a__    ,      a__				) and column:
            if all(
                isinstance(a__    ,      torch.Tensor				) and x.shape == column[0].shape and x.dtype == column[0].dtype
                for x in column				):
               return torch.stack(a__				)
         return column


      def        a			(self	:				Optional[Any]    ,      a__	:				str				):





         """simple docstring"""



         import torch

         if isinstance(a__    ,      (str, bytes, type(a__				))				):
            return value
         elif isinstance(a__    ,      (np.character, np.ndarray)				) and np.issubdtype(value.dtype    ,      np.character				):
            return value.tolist()

         __snake_case								=       {}

         if isinstance(a__    ,      (np.number, np.ndarray)				) and np.issubdtype(value.dtype    ,      np.integer				):
            __snake_case								=       {'''dtype''': torch.intaa}
         elif isinstance(a__    ,      (np.number, np.ndarray)				) and np.issubdtype(value.dtype    ,      np.floating				):
            __snake_case								=       {'''dtype''': torch.floataa}
         elif config.PIL_AVAILABLE and "PIL" in sys.modules:
            import PIL.Image

            if isinstance(a__    ,      PIL.Image.Image				):
               __snake_case								=       np.asarray(a__				)
         return torch.tensor(a__    ,      **{**default_dtype, **self.torch_tensor_kwargs}				)


      def        a			(self	:				Optional[int]    ,      a__	:				str				):





         """simple docstring"""



         import torch

         # support for torch, tf, jax etc.
         if hasattr(a__    ,      '''__array__'''				) and not isinstance(a__    ,      torch.Tensor				):
            __snake_case								=       data_struct.__array__()
         # support for nested types like struct of list of struct
         if isinstance(a__    ,      np.ndarray				):
            if data_struct.dtype == object:  # torch tensors cannot be instantied from an array of objects
               return self._consolidate([self.recursive_tensorize(a__				) for substruct in data_struct]				)
         elif isinstance(a__    ,      (list, tuple)				):
            return self._consolidate([self.recursive_tensorize(a__				) for substruct in data_struct]				)
         return self._tensorize(a__				)


      def        a			(self	:				Optional[Any]    ,      a__	:				dict				):





         """simple docstring"""



         return map_nested(self._recursive_tensorize    ,      a__    ,      map_list=a__				)


      def        a			(self	:				Optional[int]    ,      a__	:				pa.Table				):





         """simple docstring"""



         __snake_case								=       self.numpy_arrow_extractor().extract_row(a__				)
         __snake_case								=       self.python_features_decoder.decode_row(a__				)
         return self.recursive_tensorize(a__				)


      def        a			(self	:				List[Any]    ,      a__	:				pa.Table				):





         """simple docstring"""



         __snake_case								=       self.numpy_arrow_extractor().extract_column(a__				)
         __snake_case								=       self.python_features_decoder.decode_column(a__    ,      pa_table.column_names[0]				)
         __snake_case								=       self.recursive_tensorize(a__				)
         __snake_case								=       self._consolidate(a__				)
         return column


      def        a			(self	:				Optional[int]    ,      a__	:				pa.Table				):





         """simple docstring"""



         __snake_case								=       self.numpy_arrow_extractor().extract_batch(a__				)
         __snake_case								=       self.python_features_decoder.decode_batch(a__				)
         __snake_case								=       self.recursive_tensorize(a__				)
         for column_name in batch:
            __snake_case								=       self._consolidate(batch[column_name]				)
         return batch

