


"""simple docstring"""

from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)





@flax.struct.dataclass
class  lowerCamelCase__    (		_a     ):
      a			:  jnp.ndarray
      a			:  jnp.ndarray






class  lowerCamelCase__    (		nn.Module     ):
      a			:  int
      a			:  Tuple[int]           =		(16, 32, 96, 2_56)
      a			:  jnp.dtype           =		jnp.floataa
      def SCREAMING_SNAKE_CASE_						(		self						:						List[str]				):






       '''simple docstring'''


       __lowercase						=						nn.Conv(
           self.block_out_channels[0]   ,							kernel_size=(3, 3)   ,							padding=((1, 1), (1, 1))   ,							dtype=self.dtype   ,							)

       __lowercase						=						[]
       for i in range(len(self.block_out_channels				) - 1				):
        __lowercase						=						self.block_out_channels[i]
        __lowercase						=						self.block_out_channels[i + 1]
        __lowercase						=						nn.Conv(
            A_   ,							kernel_size=(3, 3)   ,							padding=((1, 1), (1, 1))   ,							dtype=self.dtype   ,							)
        blocks.append(A_				)
        __lowercase						=						nn.Conv(
            A_   ,							kernel_size=(3, 3)   ,							strides=(2, 2)   ,							padding=((1, 1), (1, 1))   ,							dtype=self.dtype   ,							)
        blocks.append(A_				)
       __lowercase						=						blocks

       __lowercase						=						nn.Conv(
           self.conditioning_embedding_channels   ,							kernel_size=(3, 3)   ,							padding=((1, 1), (1, 1))   ,							kernel_init=nn.initializers.zeros_init()   ,							bias_init=nn.initializers.zeros_init()   ,							dtype=self.dtype   ,							)






      def __call__(		self						:						Tuple   ,							A_						:						Union[str, Any]				):






       '''simple docstring'''


       __lowercase						=						self.conv_in(A_				)
       __lowercase						=						nn.silu(A_				)

       for block in self.blocks:
        __lowercase						=						block(A_				)
        __lowercase						=						nn.silu(A_				)

       __lowercase						=						self.conv_out(A_				)

       return embedding






@flax_register_to_config
class  lowerCamelCase__    (		nn.Module      , _a      , _a     ):
      a			:  int           =		32
      a			:  int           =		4
      a			:  Tuple[str]           =		(
          "CrossAttnDownBlock2D",
          "CrossAttnDownBlock2D",
          "CrossAttnDownBlock2D",
          "DownBlock2D",
      )
      a			:  Union[bool, Tuple[bool]]           =		False
      a			:  Tuple[int]           =		(3_20, 6_40, 12_80, 12_80)
      a			:  int           =		2
      a			:  Union[int, Tuple[int]]           =		8
      a			:  Optional[Union[int, Tuple[int]]]           =		None
      a			:  int           =		12_80
      a			:  float           =		0.0
      a			:  bool           =		False
      a			:  jnp.dtype           =		jnp.floataa
      a			:  bool           =		True
      a			:  int           =		0
      a			:  str           =		"rgb"
      a			:  Tuple[int]           =		(16, 32, 96, 2_56)
      def SCREAMING_SNAKE_CASE_						(		self						:						Union[str, Any]   ,							A_						:						jax.random.KeyArray				):






       '''simple docstring'''


       __lowercase						=						(1, self.in_channels, self.sample_size, self.sample_size)
       __lowercase						=						jnp.zeros(A_   ,							dtype=jnp.floataa				)
       __lowercase						=						jnp.ones((1,)   ,							dtype=jnp.intaa				)
       __lowercase						=						jnp.zeros((1, 1, self.cross_attention_dim)   ,							dtype=jnp.floataa				)
       __lowercase						=						(1, 3, self.sample_size * 8, self.sample_size * 8)
       __lowercase						=						jnp.zeros(A_   ,							dtype=jnp.floataa				)

       __lowercase				,		__lowercase						=						jax.random.split(A_				)
       __lowercase						=						{"""params""": params_rng, """dropout""": dropout_rng}

       return self.init(A_   ,							A_   ,							A_   ,							A_   ,							A_				)["params"]
      def SCREAMING_SNAKE_CASE_						(		self						:						Any				):






       '''simple docstring'''


       __lowercase						=						self.block_out_channels
       __lowercase						=						block_out_channels[0] * 4

       # If `num_attention_heads` is not defined (which is the case for most models)
       # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
       # The reason for this behavior is to correct for incorrectly named variables that were introduced
       # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
       # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
       # which is why we correct for the naming here.
       __lowercase						=						self.num_attention_heads or self.attention_head_dim

       # input
       __lowercase						=						nn.Conv(
           block_out_channels[0]   ,							kernel_size=(3, 3)   ,							strides=(1, 1)   ,							padding=((1, 1), (1, 1))   ,							dtype=self.dtype   ,							)

       # time
       __lowercase						=						FlaxTimesteps(
           block_out_channels[0]   ,							flip_sin_to_cos=self.flip_sin_to_cos   ,							freq_shift=self.config.freq_shift				)
       __lowercase						=						FlaxTimestepEmbedding(A_   ,							dtype=self.dtype				)

       __lowercase						=						FlaxControlNetConditioningEmbedding(
           conditioning_embedding_channels=block_out_channels[0]   ,							block_out_channels=self.conditioning_embedding_out_channels   ,							)

       __lowercase						=						self.only_cross_attention
       if isinstance(A_   ,							A_				):
        __lowercase						=						(only_cross_attention,) * len(self.down_block_types				)

       if isinstance(A_   ,							A_				):
        __lowercase						=						(num_attention_heads,) * len(self.down_block_types				)

       # down
       __lowercase						=						[]
       __lowercase						=						[]

       __lowercase						=						block_out_channels[0]

       __lowercase						=						nn.Conv(
           A_   ,							kernel_size=(1, 1)   ,							padding="""VALID"""   ,							kernel_init=nn.initializers.zeros_init()   ,							bias_init=nn.initializers.zeros_init()   ,							dtype=self.dtype   ,							)
       controlnet_down_blocks.append(A_				)

       for i, down_block_type in enumerate(self.down_block_types				):
        __lowercase						=						output_channel
        __lowercase						=						block_out_channels[i]
        __lowercase						=						i == len(A_				) - 1

        if down_block_type == "CrossAttnDownBlock2D":
         __lowercase						=						FlaxCrossAttnDownBlockaD(
             in_channels=A_   ,							out_channels=A_   ,							dropout=self.dropout   ,							num_layers=self.layers_per_block   ,							num_attention_heads=num_attention_heads[i]   ,							add_downsample=not is_final_block   ,							use_linear_projection=self.use_linear_projection   ,							only_cross_attention=only_cross_attention[i]   ,							dtype=self.dtype   ,							)
        else:
         __lowercase						=						FlaxDownBlockaD(
             in_channels=A_   ,							out_channels=A_   ,							dropout=self.dropout   ,							num_layers=self.layers_per_block   ,							add_downsample=not is_final_block   ,							dtype=self.dtype   ,							)

        down_blocks.append(A_				)

        for _ in range(self.layers_per_block				):
         __lowercase						=						nn.Conv(
             A_   ,							kernel_size=(1, 1)   ,							padding="""VALID"""   ,							kernel_init=nn.initializers.zeros_init()   ,							bias_init=nn.initializers.zeros_init()   ,							dtype=self.dtype   ,							)
         controlnet_down_blocks.append(A_				)

        if not is_final_block:
         __lowercase						=						nn.Conv(
             A_   ,							kernel_size=(1, 1)   ,							padding="""VALID"""   ,							kernel_init=nn.initializers.zeros_init()   ,							bias_init=nn.initializers.zeros_init()   ,							dtype=self.dtype   ,							)
         controlnet_down_blocks.append(A_				)

       __lowercase						=						down_blocks
       __lowercase						=						controlnet_down_blocks

       # mid
       __lowercase						=						block_out_channels[-1]
       __lowercase						=						FlaxUNetMidBlockaDCrossAttn(
           in_channels=A_   ,							dropout=self.dropout   ,							num_attention_heads=num_attention_heads[-1]   ,							use_linear_projection=self.use_linear_projection   ,							dtype=self.dtype   ,							)

       __lowercase						=						nn.Conv(
           A_   ,							kernel_size=(1, 1)   ,							padding="""VALID"""   ,							kernel_init=nn.initializers.zeros_init()   ,							bias_init=nn.initializers.zeros_init()   ,							dtype=self.dtype   ,							)






      def __call__(		self						:						List[str]   ,							A_						:						str   ,							A_						:						Optional[int]   ,							A_						:						int   ,							A_						:						Dict   ,							A_						:						float = 1.0   ,							A_						:						bool = True   ,							A_						:						bool = False   ,							):






       '''simple docstring'''


       __lowercase						=						self.controlnet_conditioning_channel_order
       if channel_order == "bgr":
        __lowercase						=						jnp.flip(A_   ,							axis=1				)

       # 1. time
       if not isinstance(A_   ,							jnp.ndarray				):
        __lowercase						=						jnp.array([timesteps]   ,							dtype=jnp.intaa				)
       elif isinstance(A_   ,							jnp.ndarray				) and len(timesteps.shape				) == 0:
        __lowercase						=						timesteps.astype(dtype=jnp.floataa				)
        __lowercase						=						jnp.expand_dims(A_   ,							0				)

       __lowercase						=						self.time_proj(A_				)
       __lowercase						=						self.time_embedding(A_				)

       # 2. pre-process
       __lowercase						=						jnp.transpose(A_   ,							(0, 2, 3, 1)				)
       __lowercase						=						self.conv_in(A_				)

       __lowercase						=						jnp.transpose(A_   ,							(0, 2, 3, 1)				)
       __lowercase						=						self.controlnet_cond_embedding(A_				)
       sample += controlnet_cond

       # 3. down
       __lowercase						=						(sample,)
       for down_block in self.down_blocks:
        if isinstance(A_   ,							A_				):
         __lowercase				,		__lowercase						=						down_block(A_   ,							A_   ,							A_   ,							deterministic=not train				)
        else:
         __lowercase				,		__lowercase						=						down_block(A_   ,							A_   ,							deterministic=not train				)
        down_block_res_samples += res_samples

       # 4. mid
       __lowercase						=						self.mid_block(A_   ,							A_   ,							A_   ,							deterministic=not train				)

       # 5. contronet blocks
       __lowercase						=						()
       for down_block_res_sample, controlnet_block in zip(A_   ,							self.controlnet_down_blocks				):
        __lowercase						=						controlnet_block(A_				)
        controlnet_down_block_res_samples += (down_block_res_sample,)

       __lowercase						=						controlnet_down_block_res_samples

       __lowercase						=						self.controlnet_mid_block(A_				)

       # 6. scaling
       __lowercase						=						[sample * conditioning_scale for sample in down_block_res_samples]
       mid_block_res_sample *= conditioning_scale

       if not return_dict:
        return (down_block_res_samples, mid_block_res_sample)

       return FlaxControlNetOutput(
           down_block_res_samples=A_   ,							mid_block_res_sample=A_				)




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


UpperCAmelCase__					=logging.get_logger(__name__)

# General docstring
UpperCAmelCase__					="RegNetConfig"

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

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

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





class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						int   ,							A_						:						int   ,							A_						:						int   ,							A_						:						int = 3   ,							A_						:						int = 1   ,							A_						:						int = 1   ,							A_						:						Optional[str] = "relu"   ,							):






       '''simple docstring'''


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






      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple   ,							A_						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						self.convolution(A_				)
       __lowercase						=						self.normalization(A_				)
       __lowercase						=						self.activation(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						Dict   ,							A_						:						RegNetConfig				):






       '''simple docstring'''


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






      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple   ,							A_						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						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."""				)
       __lowercase						=						self.embedder(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						Optional[int]   ,							A_						:						int   ,							A_						:						int   ,							A_						:						int = 2				):






       '''simple docstring'''


       super().__init__()
       __lowercase						=						nn.Convad(A_   ,							A_   ,							kernel_size=1   ,							stride=A_   ,							bias=A_				)
       __lowercase						=						nn.BatchNormad(A_				)






      def SCREAMING_SNAKE_CASE_						(		self						:						Dict   ,							A_						:						Tensor				):






       '''simple docstring'''


       __lowercase						=						self.convolution(A_				)
       __lowercase						=						self.normalization(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						int   ,							A_						:						int   ,							A_						:						int				):






       '''simple docstring'''


       super().__init__()

       __lowercase						=						nn.AdaptiveAvgPoolad((1, 1)				)
       __lowercase						=						nn.Sequential(
           nn.Convad(A_   ,							A_   ,							kernel_size=1				)   ,							nn.ReLU()   ,							nn.Convad(A_   ,							A_   ,							kernel_size=1				)   ,							nn.Sigmoid()   ,							)






      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]   ,							A_						:						List[Any]				):






       '''simple docstring'''


       __lowercase						=						self.pooler(A_				)
       __lowercase						=						self.attention(A_				)
       __lowercase						=						hidden_state * attention
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						Any   ,							A_						:						RegNetConfig   ,							A_						:						int   ,							A_						:						int   ,							A_						:						int = 1				):






       '''simple docstring'''


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






      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]   ,							A_						:						Tuple				):






       '''simple docstring'''


       __lowercase						=						hidden_state
       __lowercase						=						self.layer(A_				)
       __lowercase						=						self.shortcut(A_				)
       hidden_state += residual
       __lowercase						=						self.activation(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						List[str]   ,							A_						:						RegNetConfig   ,							A_						:						int   ,							A_						:						int   ,							A_						:						int = 1				):






       '''simple docstring'''


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






      def SCREAMING_SNAKE_CASE_						(		self						:						Dict   ,							A_						:						Tuple				):






       '''simple docstring'''


       __lowercase						=						hidden_state
       __lowercase						=						self.layer(A_				)
       __lowercase						=						self.shortcut(A_				)
       hidden_state += residual
       __lowercase						=						self.activation(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						Optional[int]   ,							A_						:						RegNetConfig   ,							A_						:						int   ,							A_						:						int   ,							A_						:						int = 2   ,							A_						:						int = 2   ,							):






       '''simple docstring'''


       super().__init__()

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

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






      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[Any]   ,							A_						:						Union[str, Any]				):






       '''simple docstring'''


       __lowercase						=						self.layers(A_				)
       return hidden_state






class  lowerCamelCase__    (		nn.Module     ):
      def __init__(		self						:						Tuple   ,							A_						:						RegNetConfig				):






       '''simple docstring'''


       super().__init__()
       __lowercase						=						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(
               A_   ,							config.embedding_size   ,							config.hidden_sizes[0]   ,							stride=2 if config.downsample_in_first_stage else 1   ,							depth=config.depths[0]   ,							)				)
       __lowercase						=						zip(config.hidden_sizes   ,							config.hidden_sizes[1:]				)
       for (in_channels, out_channels), depth in zip(A_   ,							config.depths[1:]				):
        self.stages.append(RegNetStage(A_   ,							A_   ,							A_   ,							depth=A_				)				)






      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple   ,							A_						:						Tensor   ,							A_						:						bool = False   ,							A_						:						bool = True				):






       '''simple docstring'''


       __lowercase						=						() if output_hidden_states else None

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

        __lowercase						=						stage_module(A_				)

       if output_hidden_states:
        __lowercase						=						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=A_   ,							hidden_states=A_				)






class  lowerCamelCase__    (		_a     ):
      a			:  List[str]           =		RegNetConfig
      a			:  int           =		"""regnet"""
      a			:  Tuple           =		"""pixel_values"""
      a			:  Dict           =		True
      def SCREAMING_SNAKE_CASE_						(		self						:						int   ,							A_						:						List[Any]				):






       '''simple docstring'''


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






      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]   ,							A_						:						Dict   ,							A_						:						List[str]=False				):






       '''simple docstring'''


       if isinstance(A_   ,							A_				):
        __lowercase						=						value






UpperCAmelCase__					=R"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\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 [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

UpperCAmelCase__					=R"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\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 [`~file_utils.ModelOutput`] instead of a plain tuple.\n"





@add_start_docstrings(
    """The bare RegNet model outputting raw features without any specific head on top."""      , _a      , )
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class  lowerCamelCase__    (		_a     ):
      def __init__(		self						:						Any   ,							A_						:						Any				):






       '''simple docstring'''


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






      @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 SCREAMING_SNAKE_CASE_						(		self						:						Dict   ,							A_						:						Tensor   ,							A_						:						Optional[bool] = None   ,							A_						:						Optional[bool] = None				):






       '''simple docstring'''


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

       __lowercase						=						self.embedder(A_				)

       __lowercase						=						self.encoder(
           A_   ,							output_hidden_states=A_   ,							return_dict=A_				)

       __lowercase						=						encoder_outputs[0]

       __lowercase						=						self.pooler(A_				)

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

       return BaseModelOutputWithPoolingAndNoAttention(
           last_hidden_state=A_   ,							pooler_output=A_   ,							hidden_states=encoder_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.
    """      , _a      , )
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class  lowerCamelCase__    (		_a     ):
      def __init__(		self						:						int   ,							A_						:						int				):






       '''simple docstring'''


       super().__init__(A_				)
       __lowercase						=						config.num_labels
       __lowercase						=						RegNetModel(A_				)
       # classification head
       __lowercase						=						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(A_				)
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT   ,							output_type=A_   ,							config_class=_CONFIG_FOR_DOC   ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,							)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]   ,							A_						:						Optional[torch.FloatTensor] = None   ,							A_						:						Optional[torch.LongTensor] = None   ,							A_						:						Optional[bool] = None   ,							A_						:						Optional[bool] = None   ,							):






       '''simple docstring'''


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

       __lowercase						=						self.regnet(A_   ,							output_hidden_states=A_   ,							return_dict=A_				)

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

       __lowercase						=						self.classifier(A_				)

       __lowercase						=						None

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

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

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