





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 ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


_a						=   logging.get_logger(__name__)

# General docstring
_a						=   "ResNetConfig"

# Base docstring
_a						=   "microsoft/resnet-50"
_a						=   [1, 2_048, 7, 7]

# Image classification docstring
_a						=   "microsoft/resnet-50"
_a						=   "tiger cat"

_a						=   [
    "microsoft/resnet-50",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]






class       __A					(					nn.Module				):

			'''simple docstring'''







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



				'''simple docstring'''







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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	self.convolution(__lowerCAmelCase							)
				lowerCamelCase__              =	self.normalization(__lowerCAmelCase							)
				lowerCamelCase__              =	self.activation(__lowerCAmelCase							)
				return hidden_state






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				super().__init__()
				lowerCamelCase__              =	ResNetConvLayer(
				    config.num_channels ,						config.embedding_size ,						kernel_size=7 ,						stride=2 ,						activation=config.hidden_act							)
				lowerCamelCase__              =	nn.MaxPoolad(kernel_size=3 ,						stride=2 ,						padding=1							)
				lowerCamelCase__              =	config.num_channels





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''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							)
				lowerCamelCase__              =	self.pooler(__lowerCAmelCase							)
				return embedding






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = 2							):



				'''simple docstring'''







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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	self.convolution(__lowerCAmelCase							)
				lowerCamelCase__              =	self.normalization(__lowerCAmelCase							)
				return hidden_state






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = 1 ,						__lowerCAmelCase = "relu"							):



				'''simple docstring'''







				super().__init__()
				lowerCamelCase__              =	in_channels != out_channels or stride != 1
				lowerCamelCase__              =	(
				    ResNetShortCut(__lowerCAmelCase ,						__lowerCAmelCase ,						stride=__lowerCAmelCase							) if should_apply_shortcut else nn.Identity()
				)
				lowerCamelCase__              =	nn.Sequential(
				    ResNetConvLayer(__lowerCAmelCase ,						__lowerCAmelCase ,						stride=__lowerCAmelCase							) ,						ResNetConvLayer(__lowerCAmelCase ,						__lowerCAmelCase ,						activation=__lowerCAmelCase							) ,						)
				lowerCamelCase__              =	ACTaFN[activation]





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	hidden_state
				lowerCamelCase__              =	self.layer(__lowerCAmelCase							)
				lowerCamelCase__              =	self.shortcut(__lowerCAmelCase							)
				hidden_state += residual
				lowerCamelCase__              =	self.activation(__lowerCAmelCase							)
				return hidden_state






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = 1 ,						__lowerCAmelCase = "relu" ,						__lowerCAmelCase = 4							):



				'''simple docstring'''







				super().__init__()
				lowerCamelCase__              =	in_channels != out_channels or stride != 1
				lowerCamelCase__              =	out_channels // reduction
				lowerCamelCase__              =	(
				    ResNetShortCut(__lowerCAmelCase ,						__lowerCAmelCase ,						stride=__lowerCAmelCase							) if should_apply_shortcut else nn.Identity()
				)
				lowerCamelCase__              =	nn.Sequential(
				    ResNetConvLayer(__lowerCAmelCase ,						__lowerCAmelCase ,						kernel_size=1							) ,						ResNetConvLayer(__lowerCAmelCase ,						__lowerCAmelCase ,						stride=__lowerCAmelCase							) ,						ResNetConvLayer(__lowerCAmelCase ,						__lowerCAmelCase ,						kernel_size=1 ,						activation=__lowerCAmelCase							) ,						)
				lowerCamelCase__              =	ACTaFN[activation]





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	hidden_state
				lowerCamelCase__              =	self.layer(__lowerCAmelCase							)
				lowerCamelCase__              =	self.shortcut(__lowerCAmelCase							)
				hidden_state += residual
				lowerCamelCase__              =	self.activation(__lowerCAmelCase							)
				return hidden_state






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = 2 ,						__lowerCAmelCase = 2 ,						):



				'''simple docstring'''







				super().__init__()

				lowerCamelCase__              =	ResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else ResNetBasicLayer

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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	input
				for layer in self.layers:
					lowerCamelCase__              =	layer(__lowerCAmelCase							)
				return hidden_state






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase							):



				'''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(
				    ResNetStage(
				        __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(ResNetStage(__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						depth=__lowerCAmelCase							)							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase = False ,						__lowerCAmelCase = True							):



				'''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       __A					(					lowerCAmelCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		ResNetConfig
			lowerCAmelCase_							  =		"""resnet"""
			lowerCAmelCase_							  =		"""pixel_values"""
			lowerCAmelCase_							  =		True





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''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 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase=False							):



				'''simple docstring'''







				if isinstance(__lowerCAmelCase ,						__lowerCAmelCase							):
					lowerCamelCase__              =	value



_a						=   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 ([`ResNetConfig`]): 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"

_a						=   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 [`~utils.ModelOutput`] instead of a plain tuple.\n"






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

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				super().__init__(__lowerCAmelCase							)
				lowerCamelCase__              =	config
				lowerCamelCase__              =	ResNetEmbeddings(__lowerCAmelCase							)
				lowerCamelCase__              =	ResNetEncoder(__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 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None							):



				'''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(
    """
    ResNet 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				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				super().__init__(__lowerCAmelCase							)
				lowerCamelCase__              =	config.num_labels
				lowerCamelCase__              =	ResNetModel(__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 	__lowerCamelCase       (					self ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None ,						):



				'''simple docstring'''







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

				lowerCamelCase__              =	self.resnet(__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							)






@add_start_docstrings(
    """
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    """    ,       lowerCAmelCase    ,       )
class       __A					(					lowerCAmelCase    ,       lowerCAmelCase				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				super().__init__(__lowerCAmelCase							)
				super()._init_backbone(__lowerCAmelCase							)

				lowerCamelCase__              =	[config.embedding_size] + config.hidden_sizes
				lowerCamelCase__              =	ResNetEmbeddings(__lowerCAmelCase							)
				lowerCamelCase__              =	ResNetEncoder(__lowerCAmelCase							)

				# initialize weights and apply final processing
				self.post_init()





			@add_start_docstrings_to_model_forward(__lowerCAmelCase							)
			@replace_return_docstrings(output_type=__lowerCAmelCase ,						config_class=_CONFIG_FOR_DOC							)
			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None							):



				'''simple docstring'''







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

				lowerCamelCase__              =	self.embedder(__lowerCAmelCase							)

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

				lowerCamelCase__              =	outputs.hidden_states

				lowerCamelCase__              =	()
				for idx, stage in enumerate(self.stage_names							):
					if stage in self.out_features:
						feature_maps += (hidden_states[idx],)

				if not return_dict:
					lowerCamelCase__              =	(feature_maps,)
					if output_hidden_states:
						output += (outputs.hidden_states,)
					return output

				return BackboneOutput(
				    feature_maps=__lowerCAmelCase ,						hidden_states=outputs.hidden_states if output_hidden_states else None ,						attentions=__lowerCAmelCase ,						)







# This is the module that test_patching.py uses to test patch_submodule()

import os  # noqa: this is just for tests
import os as renamed_os  # noqa: this is just for tests
from os import path  # noqa: this is just for tests
from os import path as renamed_path  # noqa: this is just for tests
from os.path import join  # noqa: this is just for tests
from os.path import join as renamed_join  # noqa: this is just for tests


_a						=   open  # noqa: we just need to have a builtin inside this module to test it properly



