


import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map



def     __lowerCAmelCase						(							snake_case	:    Union[dict, list, tuple, torch.Tensor] )       ->							List[Tuple[int, ...]]:
   __lowerCamelCase:				Dict     =	[]
   if isinstance(snake_case						,      snake_case ):
      for v in tree.values():
         shapes.extend(_fetch_dims(snake_case ) )
   elif isinstance(snake_case						,      (list, tuple) ):
      for t in tree:
         shapes.extend(_fetch_dims(snake_case ) )
   elif isinstance(snake_case						,      torch.Tensor ):
      shapes.append(tree.shape )
   else:
      raise ValueError("""Not supported""" )

   return shapes



@torch.jit.ignore
def     __lowerCAmelCase						(							snake_case	:    int						,      snake_case	:    Tuple[int, ...] )       ->							Tuple[int, ...]:
   __lowerCamelCase:				Union[str, Any]     =	[]
   for d in reversed(snake_case ):
      idx.append(flat_idx % d )
      __lowerCamelCase:				Optional[int]     =	flat_idx // d

   return tuple(reversed(snake_case ) )



@torch.jit.ignore
def     __lowerCAmelCase						(							snake_case	:    Sequence[int]						,      snake_case	:    Sequence[int]						,      snake_case	:    Sequence[int]						,      snake_case	:    Optional[Sequence[bool]] = None						,      snake_case	:    Optional[Sequence[bool]] = None						,      )       ->							List[Tuple[slice, ...]]:

   # start_edges and end_edges both indicate whether, starting from any given
   # dimension, the start/end index is at the top/bottom edge of the
   # corresponding tensor, modeled as a tree
   def reduce_edge_list(snake_case	:    List[bool] ) -> None:
      __lowerCamelCase:				List[Any]     =	True
      for i in range(len(snake_case ) ):
         __lowerCamelCase:				str     =	-1 * (i + 1)
         l[reversed_idx] &= tally
         __lowerCamelCase:				List[Any]     =	l[reversed_idx]

   if start_edges is None:
      __lowerCamelCase:				str     =	[s == 0 for s in start]
      reduce_edge_list(snake_case )
   if end_edges is None:
      __lowerCamelCase:				Union[str, Any]     =	[e == (d - 1) for e, d in zip(snake_case						,      snake_case )]
      reduce_edge_list(snake_case )

   # Base cases. Either start/end are empty and we're done, or the final,
   # one-dimensional tensor can be simply sliced
   if len(snake_case ) == 0:
      return [()]
   elif len(snake_case ) == 1:
      return [(slice(start[0]						,      end[0] + 1 ),)]

   __lowerCamelCase:				List[Tuple[slice, ...]]     =	[]
   __lowerCamelCase:				List[slice]     =	[]

   # Dimensions common to start and end can be selected directly
   for s, e in zip(snake_case						,      snake_case ):
      if s == e:
         path_list.append(slice(snake_case						,      s + 1 ) )
      else:
         break

   __lowerCamelCase:				Tuple[slice, ...]     =	tuple(snake_case )
   __lowerCamelCase:				Dict     =	len(snake_case )

   # start == end, and we're done
   if divergence_idx == len(snake_case ):
      return [path]

   def upper() -> Tuple[Tuple[slice, ...], ...]:
      assert start_edges is not None
      assert end_edges is not None

      __lowerCamelCase:				Dict     =	start[divergence_idx]
      return tuple(
          path + (slice(snake_case						,      sdi + 1 ),) + s
          for s in _get_minimal_slice_set(
              start[divergence_idx + 1 :]						,      [d - 1 for d in dims[divergence_idx + 1 :]]						,      dims[divergence_idx + 1 :]						,      start_edges=start_edges[divergence_idx + 1 :]						,      end_edges=[True for _ in end_edges[divergence_idx + 1 :]]						,      ) )

   def lower() -> Tuple[Tuple[slice, ...], ...]:
      assert start_edges is not None
      assert end_edges is not None

      __lowerCamelCase:				List[Any]     =	end[divergence_idx]
      return tuple(
          path + (slice(snake_case						,      edi + 1 ),) + s
          for s in _get_minimal_slice_set(
              [0 for _ in start[divergence_idx + 1 :]]						,      end[divergence_idx + 1 :]						,      dims[divergence_idx + 1 :]						,      start_edges=[True for _ in start_edges[divergence_idx + 1 :]]						,      end_edges=end_edges[divergence_idx + 1 :]						,      ) )

   # If both start and end are at the edges of the subtree rooted at
   # divergence_idx, we can just select the whole subtree at once
   if start_edges[divergence_idx] and end_edges[divergence_idx]:
      slices.append(path + (slice(start[divergence_idx]						,      end[divergence_idx] + 1 ),) )
   # If just start is at the edge, we can grab almost all of the subtree,
   # treating only the ragged bottom edge as an edge case
   elif start_edges[divergence_idx]:
      slices.append(path + (slice(start[divergence_idx]						,      end[divergence_idx] ),) )
      slices.extend(lower() )
   # Analogous to the previous case, but the top is ragged this time
   elif end_edges[divergence_idx]:
      slices.extend(upper() )
      slices.append(path + (slice(start[divergence_idx] + 1						,      end[divergence_idx] + 1 ),) )
   # If both sides of the range are ragged, we need to handle both sides
   # separately. If there's contiguous meat in between them, we can index it
   # in one big chunk
   else:
      slices.extend(upper() )
      __lowerCamelCase:				List[str]     =	end[divergence_idx] - start[divergence_idx]
      if middle_ground > 1:
         slices.append(path + (slice(start[divergence_idx] + 1						,      end[divergence_idx] ),) )
      slices.extend(lower() )

   return slices



@torch.jit.ignore
def     __lowerCAmelCase						(							snake_case	:    torch.Tensor						,      snake_case	:    int						,      snake_case	:    int						,      snake_case	:    int )       ->							torch.Tensor:
   __lowerCamelCase:				Any     =	t.shape[:no_batch_dims]
   __lowerCamelCase:				Dict     =	list(_flat_idx_to_idx(snake_case						,      snake_case ) )
   # _get_minimal_slice_set is inclusive
   __lowerCamelCase:				Optional[int]     =	list(_flat_idx_to_idx(flat_end - 1						,      snake_case ) )

   # Get an ordered list of slices to perform
   __lowerCamelCase:				str     =	_get_minimal_slice_set(
       snake_case						,      snake_case						,      snake_case						,      )

   __lowerCamelCase:				Tuple     =	[t[s] for s in slices]

   return torch.cat([s.view((-1,) + t.shape[no_batch_dims:] ) for s in sliced_tensors] )



def     __lowerCAmelCase						(							snake_case	:    Callable						,      snake_case	:    Dict[str, Any]						,      snake_case	:    int						,      snake_case	:    int						,      snake_case	:    bool = False						,      snake_case	:    Any = None						,      snake_case	:    bool = False						,      )       ->							Any:
   if not (len(snake_case ) > 0):
      raise ValueError("""Must provide at least one input""" )

   __lowerCamelCase:				Optional[int]     =	[shape[:no_batch_dims] for shape in _fetch_dims(snake_case )]
   __lowerCamelCase:				int     =	tuple([max(snake_case ) for s in zip(*snake_case )] )

   def _prep_inputs(snake_case	:    torch.Tensor ) -> torch.Tensor:
      if not low_mem:
         if not sum(t.shape[:no_batch_dims] ) == no_batch_dims:
            __lowerCamelCase:				Optional[Any]     =	t.expand(orig_batch_dims + t.shape[no_batch_dims:] )
         __lowerCamelCase:				Union[str, Any]     =	t.reshape(-1						,      *t.shape[no_batch_dims:] )
      else:
         __lowerCamelCase:				Union[str, Any]     =	t.expand(orig_batch_dims + t.shape[no_batch_dims:] )
      return t

   __lowerCamelCase:				Dict[str, Any]     =	tensor_tree_map(_prep_inputs						,      snake_case )
   __lowerCamelCase:				Optional[Any]     =	None
   if _out is not None:
      __lowerCamelCase:				List[Any]     =	tensor_tree_map(lambda snake_case : t.view([-1] + list(t.shape[no_batch_dims:] ) )						,      _out )

   __lowerCamelCase:				List[Any]     =	1
   for d in orig_batch_dims:
      flat_batch_dim *= d

   __lowerCamelCase:				int     =	flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

   def _select_chunk(snake_case	:    torch.Tensor ) -> torch.Tensor:
      return t[i : i + chunk_size] if t.shape[0] != 1 else t

   __lowerCamelCase:				List[Any]     =	0
   __lowerCamelCase:				Optional[Any]     =	prepped_outputs
   for _ in range(snake_case ):
      # Chunk the input
      if not low_mem:
         __lowerCamelCase:				int     =	_select_chunk
      else:
         __lowerCamelCase:				Optional[int]     =	partial(
             _chunk_slice						,      flat_start=snake_case						,      flat_end=min(snake_case						,      i + chunk_size )						,      no_batch_dims=len(snake_case )						,      )

      __lowerCamelCase:				Dict[str, Any]     =	tensor_tree_map(snake_case						,      snake_case )

      # Run the layer on the chunk
      __lowerCamelCase:				Union[str, Any]     =	layer(**snake_case )

      # Allocate space for the output
      if out is None:
         __lowerCamelCase:				Dict     =	tensor_tree_map(lambda snake_case : t.new_zeros((flat_batch_dim,) + t.shape[1:] )						,      snake_case )

      # Put the chunk in its pre-allocated space
      if isinstance(snake_case						,      snake_case ):

         def assign(snake_case	:    dict						,      snake_case	:    dict ) -> None:
            for k, v in da.items():
               if isinstance(snake_case						,      snake_case ):
                  assign(snake_case						,      da[k] )
               else:
                  if _add_into_out:
                     v[i : i + chunk_size] += da[k]
                  else:
                     __lowerCamelCase:				List[str]     =	da[k]

         assign(snake_case						,      snake_case )
      elif isinstance(snake_case						,      snake_case ):
         for xa, xa in zip(snake_case						,      snake_case ):
            if _add_into_out:
               xa[i : i + chunk_size] += xa
            else:
               __lowerCamelCase:				Optional[int]     =	xa
      elif isinstance(snake_case						,      torch.Tensor ):
         if _add_into_out:
            out[i : i + chunk_size] += output_chunk
         else:
            __lowerCamelCase:				Optional[Any]     =	output_chunk
      else:
         raise ValueError("""Not supported""" )

      i += chunk_size

   __lowerCamelCase:				Union[str, Any]     =	tensor_tree_map(lambda snake_case : t.view(orig_batch_dims + t.shape[1:] )						,      snake_case )

   return out







class        a					:





     def __init__(    self					:				Dict    ,	SCREAMING_SNAKE_CASE_					:				int = 512    ,	):
        __lowerCamelCase:				Any     =	max_chunk_size
        __lowerCamelCase:				Optional[int]     =	None
        __lowerCamelCase:				Optional[tuple]     =	None





     def  SCREAMING_SNAKE_CASE__       (    self					:				Tuple    ,	SCREAMING_SNAKE_CASE_					:				Callable    ,	SCREAMING_SNAKE_CASE_					:				tuple    ,	SCREAMING_SNAKE_CASE_					:				int							):
        logging.info("""Tuning chunk size..."""							)

        if min_chunk_size >= self.max_chunk_size:
           return min_chunk_size

        __lowerCamelCase:				List[int]     =	[2**l for l in range(int(math.log(self.max_chunk_size    ,	2							)							) + 1							)]
        __lowerCamelCase:				Optional[Any]     =	[c for c in candidates if c > min_chunk_size]
        __lowerCamelCase:				Union[str, Any]     =	[min_chunk_size] + candidates
        candidates[-1] += 4

        def test_chunk_size(SCREAMING_SNAKE_CASE_					:				int							) -> bool:
           try:
              with torch.no_grad():
                 fn(*SCREAMING_SNAKE_CASE_    ,	chunk_size=SCREAMING_SNAKE_CASE_							)
              return True
           except RuntimeError:
              return False

        __lowerCamelCase:				List[str]     =	0
        __lowerCamelCase:				Any     =	len(SCREAMING_SNAKE_CASE_							) - 1
        while i > min_viable_chunk_size_index:
           __lowerCamelCase:				List[str]     =	test_chunk_size(candidates[i]							)
           if not viable:
              __lowerCamelCase:				Optional[int]     =	(min_viable_chunk_size_index + i) // 2
           else:
              __lowerCamelCase:				Optional[Any]     =	i
              __lowerCamelCase:				List[str]     =	(i + len(SCREAMING_SNAKE_CASE_							) - 1) // 2

        return candidates[min_viable_chunk_size_index]





     def  SCREAMING_SNAKE_CASE__       (    self					:				int    ,	SCREAMING_SNAKE_CASE_					:				Iterable    ,	SCREAMING_SNAKE_CASE_					:				Iterable							):
        __lowerCamelCase:				int     =	True
        for aa, aa in zip(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							):
           assert type(SCREAMING_SNAKE_CASE_							) == type(SCREAMING_SNAKE_CASE_							)
           if isinstance(SCREAMING_SNAKE_CASE_    ,	(list, tuple)							):
              consistent &= self._compare_arg_caches(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							)
           elif isinstance(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							):
              __lowerCamelCase:				int     =	[v for _, v in sorted(aa.items()    ,	key=lambda SCREAMING_SNAKE_CASE_							: x[0]							)]
              __lowerCamelCase:				Dict     =	[v for _, v in sorted(aa.items()    ,	key=lambda SCREAMING_SNAKE_CASE_							: x[0]							)]
              consistent &= self._compare_arg_caches(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							)
           else:
              consistent &= aa == aa

        return consistent







     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[int]    ,	SCREAMING_SNAKE_CASE_					:				Callable    ,	SCREAMING_SNAKE_CASE_					:				tuple    ,	SCREAMING_SNAKE_CASE_					:				int    ,	):
        __lowerCamelCase:				int     =	True
        __lowerCamelCase:				tuple     =	tree_map(lambda SCREAMING_SNAKE_CASE_							: a.shape if isinstance(SCREAMING_SNAKE_CASE_    ,	torch.Tensor							) else a    ,	SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							)
        if self.cached_arg_data is not None:
           # If args have changed shape/value, we need to re-tune
           assert len(self.cached_arg_data							) == len(SCREAMING_SNAKE_CASE_							)
           __lowerCamelCase:				List[str]     =	self._compare_arg_caches(self.cached_arg_data    ,	SCREAMING_SNAKE_CASE_							)
        else:
           # Otherwise, we can reuse the precomputed value
           __lowerCamelCase:				Optional[int]     =	False

        if not consistent:
           __lowerCamelCase:				str     =	self._determine_favorable_chunk_size(
               SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	)
           __lowerCamelCase:				Optional[int]     =	arg_data

        assert self.cached_chunk_size is not None

        return self.cached_chunk_size






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		:				Optional[Any]									=       logging.get_logger(__name__)

# General docstring
_A		:				int									=       '''ResNetConfig'''

# Base docstring
_A		:				Union[str, Any]									=       '''microsoft/resnet-50'''
_A		:				str									=       [1, 2_048, 7, 7]

# Image classification docstring
_A		:				List[str]									=       '''microsoft/resnet-50'''
_A		:				Optional[int]									=       '''tiger cat'''

_A		:				Optional[Any]									=       [
    '''microsoft/resnet-50''',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]



class        a					(			nn.Module   ):





     def __init__(    self					:				str    ,	SCREAMING_SNAKE_CASE_					:				int    ,	SCREAMING_SNAKE_CASE_					:				int    ,	SCREAMING_SNAKE_CASE_					:				int = 3    ,	SCREAMING_SNAKE_CASE_					:				int = 1    ,	SCREAMING_SNAKE_CASE_					:				str = "relu"							):
        super().__init__()
        __lowerCamelCase:				Optional[int]     =	nn.Convad(
            SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	kernel_size=SCREAMING_SNAKE_CASE_    ,	stride=SCREAMING_SNAKE_CASE_    ,	padding=kernel_size // 2    ,	bias=SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				List[str]     =	nn.BatchNormad(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Dict     =	ACTaFN[activation] if activation is not None else nn.Identity()







     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[int]    ,	SCREAMING_SNAKE_CASE_					:				Tensor							):
        __lowerCamelCase:				Optional[int]     =	self.convolution(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Tuple     =	self.normalization(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Union[str, Any]     =	self.activation(SCREAMING_SNAKE_CASE_							)
        return hidden_state



class        a					(			nn.Module   ):





     def __init__(    self					:				Optional[Any]    ,	SCREAMING_SNAKE_CASE_					:				ResNetConfig							):
        super().__init__()
        __lowerCamelCase:				Tuple     =	ResNetConvLayer(
            config.num_channels    ,	config.embedding_size    ,	kernel_size=7    ,	stride=2    ,	activation=config.hidden_act							)
        __lowerCamelCase:				Tuple     =	nn.MaxPoolad(kernel_size=3    ,	stride=2    ,	padding=1							)
        __lowerCamelCase:				List[Any]     =	config.num_channels







     def  SCREAMING_SNAKE_CASE__       (    self					:				Dict    ,	SCREAMING_SNAKE_CASE_					:				Tensor							):
        __lowerCamelCase:				Optional[Any]     =	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:				Optional[Any]     =	self.embedder(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Tuple     =	self.pooler(SCREAMING_SNAKE_CASE_							)
        return embedding



class        a					(			nn.Module   ):





     def __init__(    self					:				int    ,	SCREAMING_SNAKE_CASE_					:				int    ,	SCREAMING_SNAKE_CASE_					:				int    ,	SCREAMING_SNAKE_CASE_					:				int = 2							):
        super().__init__()
        __lowerCamelCase:				Optional[Any]     =	nn.Convad(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	kernel_size=1    ,	stride=SCREAMING_SNAKE_CASE_    ,	bias=SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Optional[Any]     =	nn.BatchNormad(SCREAMING_SNAKE_CASE_							)







     def  SCREAMING_SNAKE_CASE__       (    self					:				int    ,	SCREAMING_SNAKE_CASE_					:				Tensor							):
        __lowerCamelCase:				Optional[int]     =	self.convolution(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				List[Any]     =	self.normalization(SCREAMING_SNAKE_CASE_							)
        return hidden_state



class        a					(			nn.Module   ):





     def __init__(    self					:				Union[str, Any]    ,	SCREAMING_SNAKE_CASE_					:				int    ,	SCREAMING_SNAKE_CASE_					:				int    ,	SCREAMING_SNAKE_CASE_					:				int = 1    ,	SCREAMING_SNAKE_CASE_					:				str = "relu"							):
        super().__init__()
        __lowerCamelCase:				Union[str, Any]     =	in_channels != out_channels or stride != 1
        __lowerCamelCase:				int     =	(
            ResNetShortCut(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	stride=SCREAMING_SNAKE_CASE_							) if should_apply_shortcut else nn.Identity()
        )
        __lowerCamelCase:				List[str]     =	nn.Sequential(
            ResNetConvLayer(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	stride=SCREAMING_SNAKE_CASE_							)    ,	ResNetConvLayer(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	activation=SCREAMING_SNAKE_CASE_							)    ,	)
        __lowerCamelCase:				Dict     =	ACTaFN[activation]







     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[int]    ,	SCREAMING_SNAKE_CASE_					:				Any							):
        __lowerCamelCase:				List[Any]     =	hidden_state
        __lowerCamelCase:				Optional[Any]     =	self.layer(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Union[str, Any]     =	self.shortcut(SCREAMING_SNAKE_CASE_							)
        hidden_state += residual
        __lowerCamelCase:				str     =	self.activation(SCREAMING_SNAKE_CASE_							)
        return hidden_state



class        a					(			nn.Module   ):





     def __init__(    self					:				Union[str, Any]    ,	SCREAMING_SNAKE_CASE_					:				int    ,	SCREAMING_SNAKE_CASE_					:				int    ,	SCREAMING_SNAKE_CASE_					:				int = 1    ,	SCREAMING_SNAKE_CASE_					:				str = "relu"    ,	SCREAMING_SNAKE_CASE_					:				int = 4							):
        super().__init__()
        __lowerCamelCase:				str     =	in_channels != out_channels or stride != 1
        __lowerCamelCase:				Any     =	out_channels // reduction
        __lowerCamelCase:				str     =	(
            ResNetShortCut(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	stride=SCREAMING_SNAKE_CASE_							) if should_apply_shortcut else nn.Identity()
        )
        __lowerCamelCase:				Any     =	nn.Sequential(
            ResNetConvLayer(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	kernel_size=1							)    ,	ResNetConvLayer(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	stride=SCREAMING_SNAKE_CASE_							)    ,	ResNetConvLayer(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	kernel_size=1    ,	activation=SCREAMING_SNAKE_CASE_							)    ,	)
        __lowerCamelCase:				str     =	ACTaFN[activation]







     def  SCREAMING_SNAKE_CASE__       (    self					:				int    ,	SCREAMING_SNAKE_CASE_					:				Optional[Any]							):
        __lowerCamelCase:				Optional[Any]     =	hidden_state
        __lowerCamelCase:				int     =	self.layer(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Dict     =	self.shortcut(SCREAMING_SNAKE_CASE_							)
        hidden_state += residual
        __lowerCamelCase:				List[str]     =	self.activation(SCREAMING_SNAKE_CASE_							)
        return hidden_state



class        a					(			nn.Module   ):





     def __init__(    self					:				Union[str, Any]    ,	SCREAMING_SNAKE_CASE_					:				ResNetConfig    ,	SCREAMING_SNAKE_CASE_					:				int    ,	SCREAMING_SNAKE_CASE_					:				int    ,	SCREAMING_SNAKE_CASE_					:				int = 2    ,	SCREAMING_SNAKE_CASE_					:				int = 2    ,	):
        super().__init__()

        __lowerCamelCase:				str     =	ResNetBottleNeckLayer if config.layer_type == """bottleneck""" else ResNetBasicLayer

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







     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]    ,	SCREAMING_SNAKE_CASE_					:				Tensor							):
        __lowerCamelCase:				Tuple     =	input
        for layer in self.layers:
           __lowerCamelCase:				int     =	layer(SCREAMING_SNAKE_CASE_							)
        return hidden_state



class        a					(			nn.Module   ):





     def __init__(    self					:				Tuple    ,	SCREAMING_SNAKE_CASE_					:				ResNetConfig							):
        super().__init__()
        __lowerCamelCase:				Any     =	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(
                SCREAMING_SNAKE_CASE_    ,	config.embedding_size    ,	config.hidden_sizes[0]    ,	stride=2 if config.downsample_in_first_stage else 1    ,	depth=config.depths[0]    ,	)							)
        __lowerCamelCase:				List[str]     =	zip(config.hidden_sizes    ,	config.hidden_sizes[1:]							)
        for (in_channels, out_channels), depth in zip(SCREAMING_SNAKE_CASE_    ,	config.depths[1:]							):
           self.stages.append(ResNetStage(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	depth=SCREAMING_SNAKE_CASE_							)							)







     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[int]    ,	SCREAMING_SNAKE_CASE_					:				Tensor    ,	SCREAMING_SNAKE_CASE_					:				bool = False    ,	SCREAMING_SNAKE_CASE_					:				bool = True							):
        __lowerCamelCase:				List[Any]     =	() if output_hidden_states else None

        for stage_module in self.stages:
           if output_hidden_states:
              __lowerCamelCase:				Optional[Any]     =	hidden_states + (hidden_state,)

           __lowerCamelCase:				Union[str, Any]     =	stage_module(SCREAMING_SNAKE_CASE_							)

        if output_hidden_states:
           __lowerCamelCase:				Union[str, Any]     =	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=SCREAMING_SNAKE_CASE_    ,	hidden_states=SCREAMING_SNAKE_CASE_    ,	)



class        a					(			_UpperCAmelCase   ):
     UpperCAmelCase__ :  Dict						      =	ResNetConfig
     UpperCAmelCase__ :  str						      =	"resnet"
     UpperCAmelCase__ :  Optional[int]						      =	"pixel_values"
     UpperCAmelCase__ :  List[str]						      =	True





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







     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]    ,	SCREAMING_SNAKE_CASE_					:				str    ,	SCREAMING_SNAKE_CASE_					:				str=False							):
        if isinstance(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_							):
           __lowerCamelCase:				Union[str, Any]     =	value



_A		:				Dict									=       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."     ,_UpperCAmelCase     ,)
class        a					(			_UpperCAmelCase   ):





     def __init__(    self					:				List[Any]    ,	SCREAMING_SNAKE_CASE_					:				List[Any]							):
        super().__init__(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Dict     =	config
        __lowerCamelCase:				Tuple     =	ResNetEmbeddings(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				List[str]     =	ResNetEncoder(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				List[Any]     =	nn.AdaptiveAvgPoolad((1, 1)							)
        # Initialize weights and apply final processing
        self.post_init()







     @add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_							)
     @add_code_sample_docstrings(
         checkpoint=_CHECKPOINT_FOR_DOC    ,	output_type=SCREAMING_SNAKE_CASE_    ,	config_class=_CONFIG_FOR_DOC    ,	modality="""vision"""    ,	expected_output=_EXPECTED_OUTPUT_SHAPE    ,	)
     def  SCREAMING_SNAKE_CASE__       (    self					:				Dict    ,	SCREAMING_SNAKE_CASE_					:				Tensor    ,	SCREAMING_SNAKE_CASE_					:				Optional[bool] = None    ,	SCREAMING_SNAKE_CASE_					:				Optional[bool] = None							):
        __lowerCamelCase:				Dict     =	(
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        __lowerCamelCase:				List[Any]     =	return_dict if return_dict is not None else self.config.use_return_dict

        __lowerCamelCase:				Any     =	self.embedder(SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				List[str]     =	self.encoder(
            SCREAMING_SNAKE_CASE_    ,	output_hidden_states=SCREAMING_SNAKE_CASE_    ,	return_dict=SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				List[str]     =	encoder_outputs[0]

        __lowerCamelCase:				List[Any]     =	self.pooler(SCREAMING_SNAKE_CASE_							)

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

        return BaseModelOutputWithPoolingAndNoAttention(
            last_hidden_state=SCREAMING_SNAKE_CASE_    ,	pooler_output=SCREAMING_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    "     ,_UpperCAmelCase     ,)
class        a					(			_UpperCAmelCase   ):





     def __init__(    self					:				Union[str, Any]    ,	SCREAMING_SNAKE_CASE_					:				Dict							):
        super().__init__(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Optional[Any]     =	config.num_labels
        __lowerCamelCase:				int     =	ResNetModel(SCREAMING_SNAKE_CASE_							)
        # classification head
        __lowerCamelCase:				str     =	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(SCREAMING_SNAKE_CASE_							)
     @add_code_sample_docstrings(
         checkpoint=_IMAGE_CLASS_CHECKPOINT    ,	output_type=SCREAMING_SNAKE_CASE_    ,	config_class=_CONFIG_FOR_DOC    ,	expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT    ,	)
     def  SCREAMING_SNAKE_CASE__       (    self					:				str    ,	SCREAMING_SNAKE_CASE_					:				Optional[torch.FloatTensor] = None    ,	SCREAMING_SNAKE_CASE_					:				Optional[torch.LongTensor] = None    ,	SCREAMING_SNAKE_CASE_					:				Optional[bool] = None    ,	SCREAMING_SNAKE_CASE_					:				Optional[bool] = None    ,	):
        __lowerCamelCase:				Tuple     =	return_dict if return_dict is not None else self.config.use_return_dict

        __lowerCamelCase:				List[str]     =	self.resnet(SCREAMING_SNAKE_CASE_    ,	output_hidden_states=SCREAMING_SNAKE_CASE_    ,	return_dict=SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				Tuple     =	outputs.pooler_output if return_dict else outputs[1]

        __lowerCamelCase:				str     =	self.classifier(SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				Dict     =	None

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

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

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



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





     def __init__(    self					:				str    ,	SCREAMING_SNAKE_CASE_					:				Optional[int]							):
        super().__init__(SCREAMING_SNAKE_CASE_							)
        super()._init_backbone(SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				Optional[int]     =	[config.embedding_size] + config.hidden_sizes
        __lowerCamelCase:				str     =	ResNetEmbeddings(SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				Optional[Any]     =	ResNetEncoder(SCREAMING_SNAKE_CASE_							)

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







     @add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_							)
     @replace_return_docstrings(output_type=SCREAMING_SNAKE_CASE_    ,	config_class=_CONFIG_FOR_DOC							)
     def  SCREAMING_SNAKE_CASE__       (    self					:				List[str]    ,	SCREAMING_SNAKE_CASE_					:				Tensor    ,	SCREAMING_SNAKE_CASE_					:				Optional[bool] = None    ,	SCREAMING_SNAKE_CASE_					:				Optional[bool] = None							):
        __lowerCamelCase:				Dict     =	return_dict if return_dict is not None else self.config.use_return_dict
        __lowerCamelCase:				Dict     =	(
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )

        __lowerCamelCase:				str     =	self.embedder(SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				Dict     =	self.encoder(SCREAMING_SNAKE_CASE_    ,	output_hidden_states=SCREAMING_SNAKE_CASE_    ,	return_dict=SCREAMING_SNAKE_CASE_							)

        __lowerCamelCase:				List[str]     =	outputs.hidden_states

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

        if not return_dict:
           __lowerCamelCase:				Union[str, Any]     =	(feature_maps,)
           if output_hidden_states:
              output += (outputs.hidden_states,)
           return output

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


