






'''simple docstring'''





from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import logging
from .configuration_regnet import RegNetConfig


__A							:							Dict  							=		logging.get_logger(__name__)

# General docstring
__A							:							Optional[Any]  							=		"""RegNetConfig"""

# Base docstring
__A							:							Dict  							=		"""facebook/regnet-y-040"""
__A							:							int  							=		[1, 10_88, 7, 7]

# Image classification docstring
__A							:							str  							=		"""facebook/regnet-y-040"""
__A							:							Union[str, Any]  							=		"""tabby, tabby cat"""

__A							:							int  							=		[
    """facebook/regnet-y-040""",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     List[str]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     int = 3				,     __lowerCamelCase     :     int = 1				,     __lowerCamelCase     :     int = 1				,     __lowerCamelCase     :     Optional[str] = "relu"				,     )							->							int:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.Convad(
         __lowerCamelCase				,     __lowerCamelCase				,     kernel_size=__lowerCamelCase				,     stride=__lowerCamelCase				,     padding=kernel_size // 2				,     groups=__lowerCamelCase				,     bias=__lowerCamelCase				,     )
     lowerCamelCase__  =     nn.BatchNormad(__lowerCamelCase     )
     lowerCamelCase__  =     ACTaFN[activation] if activation is not None else nn.Identity()




    def    a__				(					self     :     Optional[int]				,     __lowerCamelCase     :     str     )							->							str:

     '''simple docstring'''




     lowerCamelCase__  =     self.convolution(__lowerCamelCase     )
     lowerCamelCase__  =     self.normalization(__lowerCamelCase     )
     lowerCamelCase__  =     self.activation(__lowerCamelCase     )
     return hidden_state
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Dict				,     __lowerCamelCase     :     RegNetConfig     )							->							Dict:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     RegNetConvLayer(
         config.num_channels				,     config.embedding_size				,     kernel_size=3				,     stride=2				,     activation=config.hidden_act     )
     lowerCamelCase__  =     config.num_channels




    def    a__				(					self     :     Tuple				,     __lowerCamelCase     :     Optional[int]     )							->							Optional[Any]:

     '''simple docstring'''




     lowerCamelCase__  =     pixel_values.shape[1]
     if 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."     )
     lowerCamelCase__  =     self.embedder(__lowerCamelCase     )
     return hidden_state
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     int				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     int = 2     )							->							Tuple:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.Convad(__lowerCamelCase				,     __lowerCamelCase				,     kernel_size=1				,     stride=__lowerCamelCase				,     bias=__lowerCamelCase     )
     lowerCamelCase__  =     nn.BatchNormad(__lowerCamelCase     )




    def    a__				(					self     :     List[Any]				,     __lowerCamelCase     :     Tensor     )							->							Tensor:

     '''simple docstring'''




     lowerCamelCase__  =     self.convolution(__lowerCamelCase     )
     lowerCamelCase__  =     self.normalization(__lowerCamelCase     )
     return hidden_state
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     str				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     int     )							->							str:

     '''simple docstring'''




     super().__init__()

     lowerCamelCase__  =     nn.AdaptiveAvgPoolad((1, 1)     )
     lowerCamelCase__  =     nn.Sequential(
         nn.Convad(__lowerCamelCase				,     __lowerCamelCase				,     kernel_size=1     )				,     nn.ReLU()				,     nn.Convad(__lowerCamelCase				,     __lowerCamelCase				,     kernel_size=1     )				,     nn.Sigmoid()				,     )




    def    a__				(					self     :     List[Any]				,     __lowerCamelCase     :     int     )							->							Any:

     '''simple docstring'''




     lowerCamelCase__  =     self.pooler(__lowerCamelCase     )
     lowerCamelCase__  =     self.attention(__lowerCamelCase     )
     lowerCamelCase__  =     hidden_state * attention
     return hidden_state
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Union[str, Any]				,     __lowerCamelCase     :     RegNetConfig				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     int = 1     )							->							int:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     in_channels != out_channels or stride != 1
     lowerCamelCase__  =     max(1				,     out_channels // config.groups_width     )
     lowerCamelCase__  =     (
         RegNetShortCut(__lowerCamelCase				,     __lowerCamelCase				,     stride=__lowerCamelCase     ) if should_apply_shortcut else nn.Identity()
     )
     lowerCamelCase__  =     nn.Sequential(
         RegNetConvLayer(__lowerCamelCase				,     __lowerCamelCase				,     kernel_size=1				,     activation=config.hidden_act     )				,     RegNetConvLayer(__lowerCamelCase				,     __lowerCamelCase				,     stride=__lowerCamelCase				,     groups=__lowerCamelCase				,     activation=config.hidden_act     )				,     RegNetConvLayer(__lowerCamelCase				,     __lowerCamelCase				,     kernel_size=1				,     activation=__lowerCamelCase     )				,     )
     lowerCamelCase__  =     ACTaFN[config.hidden_act]




    def    a__				(					self     :     List[str]				,     __lowerCamelCase     :     List[Any]     )							->							Dict:

     '''simple docstring'''




     lowerCamelCase__  =     hidden_state
     lowerCamelCase__  =     self.layer(__lowerCamelCase     )
     lowerCamelCase__  =     self.shortcut(__lowerCamelCase     )
     hidden_state += residual
     lowerCamelCase__  =     self.activation(__lowerCamelCase     )
     return hidden_state
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Optional[Any]				,     __lowerCamelCase     :     RegNetConfig				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     int = 1     )							->							Union[str, Any]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     in_channels != out_channels or stride != 1
     lowerCamelCase__  =     max(1				,     out_channels // config.groups_width     )
     lowerCamelCase__  =     (
         RegNetShortCut(__lowerCamelCase				,     __lowerCamelCase				,     stride=__lowerCamelCase     ) if should_apply_shortcut else nn.Identity()
     )
     lowerCamelCase__  =     nn.Sequential(
         RegNetConvLayer(__lowerCamelCase				,     __lowerCamelCase				,     kernel_size=1				,     activation=config.hidden_act     )				,     RegNetConvLayer(__lowerCamelCase				,     __lowerCamelCase				,     stride=__lowerCamelCase				,     groups=__lowerCamelCase				,     activation=config.hidden_act     )				,     RegNetSELayer(__lowerCamelCase				,     reduced_channels=int(round(in_channels / 4     )     )     )				,     RegNetConvLayer(__lowerCamelCase				,     __lowerCamelCase				,     kernel_size=1				,     activation=__lowerCamelCase     )				,     )
     lowerCamelCase__  =     ACTaFN[config.hidden_act]




    def    a__				(					self     :     Tuple				,     __lowerCamelCase     :     int     )							->							str:

     '''simple docstring'''




     lowerCamelCase__  =     hidden_state
     lowerCamelCase__  =     self.layer(__lowerCamelCase     )
     lowerCamelCase__  =     self.shortcut(__lowerCamelCase     )
     hidden_state += residual
     lowerCamelCase__  =     self.activation(__lowerCamelCase     )
     return hidden_state
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Dict				,     __lowerCamelCase     :     RegNetConfig				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     int = 2				,     __lowerCamelCase     :     int = 2				,     )							->							int:

     '''simple docstring'''




     super().__init__()

     lowerCamelCase__  =     RegNetXLayer if config.layer_type == "x" else RegNetYLayer

     lowerCamelCase__  =     nn.Sequential(
         # downsampling is done in the first layer with stride of 2
         layer(
             __lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase				,     stride=__lowerCamelCase				,     )				,     *[layer(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase     ) for _ in range(depth - 1     )]				,     )




    def    a__				(					self     :     str				,     __lowerCamelCase     :     int     )							->							str:

     '''simple docstring'''




     lowerCamelCase__  =     self.layers(__lowerCamelCase     )
     return hidden_state
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     List[Any]				,     __lowerCamelCase     :     RegNetConfig     )							->							int:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.ModuleList([]     )
     # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
     self.stages.append(
         RegNetStage(
             __lowerCamelCase				,     config.embedding_size				,     config.hidden_sizes[0]				,     stride=2 if config.downsample_in_first_stage else 1				,     depth=config.depths[0]				,     )     )
     lowerCamelCase__  =     zip(config.hidden_sizes				,     config.hidden_sizes[1:]     )
     for (in_channels, out_channels), depth in zip(__lowerCamelCase				,     config.depths[1:]     ):
      self.stages.append(RegNetStage(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase				,     depth=__lowerCamelCase     )     )




    def    a__				(					self     :     int				,     __lowerCamelCase     :     Tensor				,     __lowerCamelCase     :     bool = False				,     __lowerCamelCase     :     bool = True     )							->							BaseModelOutputWithNoAttention:

     '''simple docstring'''




     lowerCamelCase__  =     () if output_hidden_states else None

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

      lowerCamelCase__  =     stage_module(__lowerCamelCase     )

     if output_hidden_states:
      lowerCamelCase__  =     hidden_states + (hidden_state,)

     if not return_dict:
      return tuple(v for v in [hidden_state, hidden_states] if v is not None     )

     return BaseModelOutputWithNoAttention(last_hidden_state=__lowerCamelCase				,     hidden_states=__lowerCamelCase     )
class 				lowercase       (							_lowerCamelCase      ):


    '''simple docstring'''

    lowerCAmelCase__      				=							RegNetConfig
    lowerCAmelCase__      				=							"regnet"
    lowerCAmelCase__      				=							"pixel_values"
    lowerCAmelCase__      				=							True

    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Any     )							->							Optional[Any]:

     '''simple docstring'''




     if isinstance(__lowerCamelCase				,     nn.Convad     ):
      nn.init.kaiming_normal_(module.weight				,     mode="fan_out"				,     nonlinearity="relu"     )
     elif isinstance(__lowerCamelCase				,     (nn.BatchNormad, nn.GroupNorm)     ):
      nn.init.constant_(module.weight				,     1     )
      nn.init.constant_(module.bias				,     0     )




    def    a__				(					self     :     Union[str, Any]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     List[str]=False     )							->							List[Any]:

     '''simple docstring'''




     if isinstance(__lowerCamelCase				,     __lowerCamelCase     ):
      lowerCamelCase__  =     value
__A							:							List[Any]  							=		R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`RegNetConfig`]): 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 [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

__A							:							Dict  							=		R"""
    Args:
        pixel_values (`torch.FloatTensor` 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 [`~file_utils.ModelOutput`] instead of a plain tuple.
"""
@add_start_docstrings(
    "The bare RegNet model outputting raw features without any specific head on top." ,       _lowerCamelCase ,       )
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class 				lowercase       (							_lowerCamelCase      ):


    '''simple docstring'''

    def __init__(					self     :     Union[str, Any]				,     __lowerCamelCase     :     int     )							->							int:

     '''simple docstring'''




     super().__init__(__lowerCamelCase     )
     lowerCamelCase__  =     config
     lowerCamelCase__  =     RegNetEmbeddings(__lowerCamelCase     )
     lowerCamelCase__  =     RegNetEncoder(__lowerCamelCase     )
     lowerCamelCase__  =     nn.AdaptiveAvgPoolad((1, 1)     )
     # Initialize weights and apply final processing
     self.post_init()




    @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				,     )
    def    a__				(					self     :     Any				,     __lowerCamelCase     :     Tensor				,     __lowerCamelCase     :     Optional[bool] = None				,     __lowerCamelCase     :     Optional[bool] = None     )							->							BaseModelOutputWithPoolingAndNoAttention:

     '''simple docstring'''




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

     lowerCamelCase__  =     self.embedder(__lowerCamelCase     )

     lowerCamelCase__  =     self.encoder(
         __lowerCamelCase				,     output_hidden_states=__lowerCamelCase				,     return_dict=__lowerCamelCase     )

     lowerCamelCase__  =     encoder_outputs[0]

     lowerCamelCase__  =     self.pooler(__lowerCamelCase     )

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

     return BaseModelOutputWithPoolingAndNoAttention(
         last_hidden_state=__lowerCamelCase				,     pooler_output=__lowerCamelCase				,     hidden_states=encoder_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    " ,       _lowerCamelCase ,       )
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class 				lowercase       (							_lowerCamelCase      ):


    '''simple docstring'''

    def __init__(					self     :     str				,     __lowerCamelCase     :     Optional[Any]     )							->							str:

     '''simple docstring'''




     super().__init__(__lowerCamelCase     )
     lowerCamelCase__  =     config.num_labels
     lowerCamelCase__  =     RegNetModel(__lowerCamelCase     )
     # classification head
     lowerCamelCase__  =     nn.Sequential(
         nn.Flatten()				,     nn.Linear(config.hidden_sizes[-1]				,     config.num_labels     ) if config.num_labels > 0 else nn.Identity()				,     )
     # initialize weights and apply final processing
     self.post_init()




    @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				,     )
    def    a__				(					self     :     Optional[Any]				,     __lowerCamelCase     :     Optional[torch.FloatTensor] = None				,     __lowerCamelCase     :     Optional[torch.LongTensor] = None				,     __lowerCamelCase     :     Optional[bool] = None				,     __lowerCamelCase     :     Optional[bool] = None				,     )							->							ImageClassifierOutputWithNoAttention:

     '''simple docstring'''




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

     lowerCamelCase__  =     self.regnet(__lowerCamelCase				,     output_hidden_states=__lowerCamelCase				,     return_dict=__lowerCamelCase     )

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

     lowerCamelCase__  =     self.classifier(__lowerCamelCase     )

     lowerCamelCase__  =     None

     if labels is not None:
      if self.config.problem_type is None:
       if self.num_labels == 1:
        lowerCamelCase__  =     "regression"
       elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
        lowerCamelCase__  =     "single_label_classification"
       else:
        lowerCamelCase__  =     "multi_label_classification"
      if self.config.problem_type == "regression":
       lowerCamelCase__  =     MSELoss()
       if self.num_labels == 1:
        lowerCamelCase__  =     loss_fct(logits.squeeze()				,     labels.squeeze()     )
       else:
        lowerCamelCase__  =     loss_fct(__lowerCamelCase				,     __lowerCamelCase     )
      elif self.config.problem_type == "single_label_classification":
       lowerCamelCase__  =     CrossEntropyLoss()
       lowerCamelCase__  =     loss_fct(logits.view(-1				,     self.num_labels     )				,     labels.view(-1     )     )
      elif self.config.problem_type == "multi_label_classification":
       lowerCamelCase__  =     BCEWithLogitsLoss()
       lowerCamelCase__  =     loss_fct(__lowerCamelCase				,     __lowerCamelCase     )

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

     return ImageClassifierOutputWithNoAttention(loss=__lowerCamelCase				,     logits=__lowerCamelCase				,     hidden_states=outputs.hidden_states     )









'''simple docstring'''





from __future__ import annotations
class 				lowercase       :


    '''simple docstring'''

    def __init__(					self     :     Optional[int]				,     __lowerCamelCase     :     int     )							->							None:

     '''simple docstring'''




     lowerCamelCase__  =     order

     # a_{0} ... a_{k}
     lowerCamelCase__  =     [1.0] + [0.0] * order
     # b_{0} ... b_{k}
     lowerCamelCase__  =     [1.0] + [0.0] * order

     # x[n-1] ... x[n-k]
     lowerCamelCase__  =     [0.0] * self.order
     # y[n-1] ... y[n-k]
     lowerCamelCase__  =     [0.0] * self.order

    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     list[float]				,     __lowerCamelCase     :     list[float]     )							->							None:

     '''simple docstring'''




     if len(__lowerCamelCase     ) < self.order:
      lowerCamelCase__  =     [1.0, *a_coeffs]

     if len(__lowerCamelCase     ) != self.order + 1:
      lowerCamelCase__  =     (
          f'''Expected a_coeffs to have {self.order + 1} elements '''
          f'''for {self.order}-order filter, got {len(__lowerCamelCase     )}'''
      )
      raise ValueError(__lowerCamelCase     )

     if len(__lowerCamelCase     ) != self.order + 1:
      lowerCamelCase__  =     (
          f'''Expected b_coeffs to have {self.order + 1} elements '''
          f'''for {self.order}-order filter, got {len(__lowerCamelCase     )}'''
      )
      raise ValueError(__lowerCamelCase     )

     lowerCamelCase__  =     a_coeffs
     lowerCamelCase__  =     b_coeffs




    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     float     )							->							float:

     '''simple docstring'''




     lowerCamelCase__  =     0.0

     # Start at index 1 and do index 0 at the end.
     for i in range(1				,     self.order + 1     ):
      result += (
          self.b_coeffs[i] * self.input_history[i - 1]
          - self.a_coeffs[i] * self.output_history[i - 1]
      )

     lowerCamelCase__  =     (result + self.b_coeffs[0] * sample) / self.a_coeffs[0]

     lowerCamelCase__  =     self.input_history[:-1]
     lowerCamelCase__  =     self.output_history[:-1]

     lowerCamelCase__  =     sample
     lowerCamelCase__  =     result

     return result

