


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			:			str												=					logging.get_logger(__name__)

# General docstring
a			:			Union[str, Any]												=					'''ResNetConfig'''

# Base docstring
a			:			Any												=					'''microsoft/resnet-50'''
a			:			Tuple												=					[1, 2048, 7, 7]

# Image classification docstring
a			:			List[Any]												=					'''microsoft/resnet-50'''
a			:			int												=					'''tiger cat'''

a			:			List[str]												=					[
    '''microsoft/resnet-50''',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]
class 						lowerCamelCase_							(					nn.Module					):



					'''simple docstring'''







					def __init__(					self			,							snake_case_			,							snake_case_			,							snake_case_ = 3			,							snake_case_ = 1			,							snake_case_ = "relu"    ) ->      Optional[Any]:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_    ) ->      Tensor:







								'''simple docstring'''

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




class 						lowerCamelCase_							(					nn.Module					):



					'''simple docstring'''







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







								'''simple docstring'''

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






					def A   (					self			,							snake_case_    ) ->      Tensor:







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




class 						lowerCamelCase_							(					nn.Module					):



					'''simple docstring'''







					def __init__(					self			,							snake_case_			,							snake_case_			,							snake_case_ = 2    ) ->      str:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_    ) ->      Tensor:







								'''simple docstring'''

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




class 						lowerCamelCase_							(					nn.Module					):



					'''simple docstring'''







					def __init__(					self			,							snake_case_			,							snake_case_			,							snake_case_ = 1			,							snake_case_ = "relu"    ) ->      int:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_    ) ->      int:







								'''simple docstring'''

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




class 						lowerCamelCase_							(					nn.Module					):



					'''simple docstring'''







					def __init__(					self			,							snake_case_			,							snake_case_			,							snake_case_ = 1			,							snake_case_ = "relu"			,							snake_case_ = 4    ) ->      List[str]:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_    ) ->      Tuple:







								'''simple docstring'''

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




class 						lowerCamelCase_							(					nn.Module					):



					'''simple docstring'''







					def __init__(					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_ = 2			,							snake_case_ = 2			,							) ->      List[Any]:







								'''simple docstring'''

								super().__init__()

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

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






					def A   (					self			,							snake_case_    ) ->      Tensor:







								'''simple docstring'''

								__lowercase			       = input
								for layer in self.layers:
											__lowercase			       = layer(snake_case_    )
								return hidden_state




class 						lowerCamelCase_							(					nn.Module					):



					'''simple docstring'''







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







								'''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(
								    ResNetStage(
								        snake_case_			,							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(snake_case_			,							config.depths[1:]    ):
											self.stages.append(ResNetStage(snake_case_			,							snake_case_			,							snake_case_			,							depth=snake_case_    )    )






					def A   (					self			,							snake_case_			,							snake_case_ = False			,							snake_case_ = True    ) ->      BaseModelOutputWithNoAttention:







								'''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(snake_case_    )

								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=snake_case_			,							hidden_states=snake_case_			,							)




class 						lowerCamelCase_							(					lowerCAmelCase__					):



					'''simple docstring'''







					__UpperCAmelCase							  =					ResNetConfig
					__UpperCAmelCase							  =					"resnet"
					__UpperCAmelCase							  =					"pixel_values"
					__UpperCAmelCase							  =					True






					def A   (					self			,							snake_case_    ) ->      Optional[Any]:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_			,							snake_case_=False    ) ->      int:







								'''simple docstring'''

								if isinstance(snake_case_			,							snake_case_    ):
											__lowercase			       = value




a			:			str												=					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 ([`ResNetConfig`]): 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			:			Optional[int]												=					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 [`~utils.ModelOutput`] instead of a plain tuple.
'''
@add_start_docstrings(
    "The bare ResNet model outputting raw features without any specific head on top."  ,  lowerCAmelCase__  ,  )
class 						lowerCamelCase_							(					lowerCAmelCase__					):



					'''simple docstring'''







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







								'''simple docstring'''

								super().__init__(snake_case_    )
								__lowercase			       = config
								__lowercase			       = ResNetEmbeddings(snake_case_    )
								__lowercase			       = ResNetEncoder(snake_case_    )
								__lowercase			       = nn.AdaptiveAvgPoolad((1, 1)    )
								# Initialize weights and apply final processing
								self.post_init()






					@add_start_docstrings_to_model_forward(snake_case_    )
					@add_code_sample_docstrings(
					    checkpoint=_CHECKPOINT_FOR_DOC			,							output_type=snake_case_			,							config_class=_CONFIG_FOR_DOC			,							modality='''vision'''			,							expected_output=_EXPECTED_OUTPUT_SHAPE			,							)
					def A   (					self			,							snake_case_			,							snake_case_ = None			,							snake_case_ = None    ) ->      BaseModelOutputWithPoolingAndNoAttention:







								'''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(snake_case_    )

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

								__lowercase			       = encoder_outputs[0]

								__lowercase			       = self.pooler(snake_case_    )

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

								return BaseModelOutputWithPoolingAndNoAttention(
								    last_hidden_state=snake_case_			,							pooler_output=snake_case_			,							hidden_states=encoder_outputs.hidden_states			,							)




@add_start_docstrings(
    "\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "  ,  lowerCAmelCase__  ,  )
class 						lowerCamelCase_							(					lowerCAmelCase__					):



					'''simple docstring'''







					def __init__(					self			,							snake_case_    ) ->      List[Any]:







								'''simple docstring'''

								super().__init__(snake_case_    )
								__lowercase			       = config.num_labels
								__lowercase			       = ResNetModel(snake_case_    )
								# 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(snake_case_    )
					@add_code_sample_docstrings(
					    checkpoint=_IMAGE_CLASS_CHECKPOINT			,							output_type=snake_case_			,							config_class=_CONFIG_FOR_DOC			,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT			,							)
					def A   (					self			,							snake_case_ = None			,							snake_case_ = None			,							snake_case_ = None			,							snake_case_ = None			,							) ->      ImageClassifierOutputWithNoAttention:







								'''simple docstring'''

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

								__lowercase			       = self.resnet(snake_case_			,							output_hidden_states=snake_case_			,							return_dict=snake_case_    )

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

								__lowercase			       = self.classifier(snake_case_    )

								__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(snake_case_			,							snake_case_    )
											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(snake_case_			,							snake_case_    )

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

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




@add_start_docstrings(
    "\n    ResNet backbone, to be used with frameworks like DETR and MaskFormer.\n    "  ,  lowerCAmelCase__  ,  )
class 						lowerCamelCase_							(					lowerCAmelCase__  ,  lowerCAmelCase__					):



					'''simple docstring'''







					def __init__(					self			,							snake_case_    ) ->      Dict:







								'''simple docstring'''

								super().__init__(snake_case_    )
								super()._init_backbone(snake_case_    )

								__lowercase			       = [config.embedding_size] + config.hidden_sizes
								__lowercase			       = ResNetEmbeddings(snake_case_    )
								__lowercase			       = ResNetEncoder(snake_case_    )

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






					@add_start_docstrings_to_model_forward(snake_case_    )
					@replace_return_docstrings(output_type=snake_case_			,							config_class=_CONFIG_FOR_DOC    )
					def A   (					self			,							snake_case_			,							snake_case_ = None			,							snake_case_ = None    ) ->      BackboneOutput:







								'''simple docstring'''

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

								__lowercase			       = self.embedder(snake_case_    )

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

								__lowercase			       = outputs.hidden_states

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

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

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






from __future__ import annotations
class 						lowerCamelCase_							:



					'''simple docstring'''







					def __init__(					self			,							snake_case_    ) ->      None:







								'''simple docstring'''

								__lowercase			       = order

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

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






					def A   (					self			,							snake_case_			,							snake_case_    ) ->      None:







								'''simple docstring'''

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

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

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

								__lowercase			       = a_coeffs
								__lowercase			       = b_coeffs






					def A   (					self			,							snake_case_    ) ->      float:







								'''simple docstring'''

								__lowercase			       = 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]
											)

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

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

								__lowercase			       = sample
								__lowercase			       = result

								return result


