





import collections.abc
from typing import Optional, Tuple, Union

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

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


_lowerCamelCase 					=      logging.get_logger(__name__)

# General docstring
_lowerCamelCase 					=      'PoolFormerConfig'

# Base docstring
_lowerCamelCase 					=      'sail/poolformer_s12'
_lowerCamelCase 					=      [1, 512, 7, 7]

# Image classification docstring
_lowerCamelCase 					=      'sail/poolformer_s12'
_lowerCamelCase 					=      'tabby, tabby cat'

_lowerCamelCase 					=      [
    'sail/poolformer_s12',
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]






def 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				Optional[Any]       , UpperCamelCase__:				float = 0.0       , UpperCamelCase__:				bool = False     ):
		if drop_prob == 0.0 or not training:
				return input
		SCREAMING_SNAKE_CASE__								=   1 - drop_prob
		SCREAMING_SNAKE_CASE__								=   (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
		SCREAMING_SNAKE_CASE__								=   keep_prob + torch.rand(UpperCamelCase__       , dtype=input.dtype       , device=input.device     )
		random_tensor.floor_()  # binarize
		SCREAMING_SNAKE_CASE__								=   input.div(UpperCamelCase__     ) * random_tensor
		return output




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Optional[Any]				,					__A     :Optional[float] = None		) -> None:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   drop_prob


							def   _snake_case	(				self     :Any				,					__A     :torch.Tensor		) -> torch.Tensor:



									"""simple docstring"""



									return drop_path(__A				,					self.drop_prob				,					self.training		)


							def   _snake_case	(				self     :Dict		) -> str:



									"""simple docstring"""



									return "p={}".format(self.drop_prob		)




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Dict				,					__A     :Optional[Any]				,					__A     :Dict				,					__A     :List[str]				,					__A     :Optional[Any]				,					__A     :Tuple				,					__A     :Optional[Any]=None		) -> Union[str, Any]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   patch_size if isinstance(__A				,					collections.abc.Iterable		) else (patch_size, patch_size)
									SCREAMING_SNAKE_CASE__								=   stride if isinstance(__A				,					collections.abc.Iterable		) else (stride, stride)
									SCREAMING_SNAKE_CASE__								=   padding if isinstance(__A				,					collections.abc.Iterable		) else (padding, padding)

									SCREAMING_SNAKE_CASE__								=   nn.Convad(__A				,					__A				,					kernel_size=__A				,					stride=__A				,					padding=__A		)
									SCREAMING_SNAKE_CASE__								=   norm_layer(__A		) if norm_layer else nn.Identity()


							def   _snake_case	(				self     :Dict				,					__A     :Union[str, Any]		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.projection(__A		)
									SCREAMING_SNAKE_CASE__								=   self.norm(__A		)
									return embeddings




class 	UpperCamelCase_							(  nn.GroupNorm     ):


							def __init__(				self     :Dict				,					__A     :Tuple				,					**__A     :Union[str, Any]		) -> Dict:



									"""simple docstring"""



									super().__init__(1				,					__A				,					**__A		)




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :List[str]				,					__A     :Optional[int]		) -> Any:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   nn.AvgPoolad(__A				,					stride=1				,					padding=pool_size // 2				,					count_include_pad=__A		)


							def   _snake_case	(				self     :Any				,					__A     :Optional[Any]		) -> Optional[Any]:



									"""simple docstring"""



									return self.pool(__A		) - hidden_states




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Optional[Any]				,					__A     :Tuple				,					__A     :Dict				,					__A     :int				,					__A     :Any		) -> str:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   nn.Convad(__A				,					__A				,					1		)
									SCREAMING_SNAKE_CASE__								=   nn.Convad(__A				,					__A				,					1		)
									SCREAMING_SNAKE_CASE__								=   PoolFormerDropPath(__A		)
									if isinstance(config.hidden_act				,					__A		):
											SCREAMING_SNAKE_CASE__								=   ACTaFN[config.hidden_act]
									else:
											SCREAMING_SNAKE_CASE__								=   config.hidden_act


							def   _snake_case	(				self     :Union[str, Any]				,					__A     :Optional[int]		) -> Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.conva(__A		)
									SCREAMING_SNAKE_CASE__								=   self.act_fn(__A		)
									SCREAMING_SNAKE_CASE__								=   self.drop(__A		)
									SCREAMING_SNAKE_CASE__								=   self.conva(__A		)
									SCREAMING_SNAKE_CASE__								=   self.drop(__A		)

									return hidden_states




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Any				,					__A     :str				,					__A     :List[str]				,					__A     :Tuple				,					__A     :Dict				,					__A     :Union[str, Any]				,					__A     :int		) -> Optional[int]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   PoolFormerPooling(__A		)
									SCREAMING_SNAKE_CASE__								=   PoolFormerOutput(__A				,					__A				,					__A				,					__A		)
									SCREAMING_SNAKE_CASE__								=   PoolFormerGroupNorm(__A		)
									SCREAMING_SNAKE_CASE__								=   PoolFormerGroupNorm(__A		)

									# Useful for training neural nets
									SCREAMING_SNAKE_CASE__								=   PoolFormerDropPath(__A		) if drop_path > 0.0 else nn.Identity()
									SCREAMING_SNAKE_CASE__								=   config.use_layer_scale
									if config.use_layer_scale:
											SCREAMING_SNAKE_CASE__								=   nn.Parameter(
											    config.layer_scale_init_value * torch.ones((__A)		)				,					requires_grad=__A		)
											SCREAMING_SNAKE_CASE__								=   nn.Parameter(
											    config.layer_scale_init_value * torch.ones((__A)		)				,					requires_grad=__A		)


							def   _snake_case	(				self     :Optional[Any]				,					__A     :Optional[int]		) -> str:



									"""simple docstring"""



									if self.use_layer_scale:
											SCREAMING_SNAKE_CASE__								=   self.pooling(self.before_norm(__A		)		)
											SCREAMING_SNAKE_CASE__								=   self.layer_scale_a.unsqueeze(-1		).unsqueeze(-1		) * pooling_output
											# First residual connection
											SCREAMING_SNAKE_CASE__								=   hidden_states + self.drop_path(__A		)
											SCREAMING_SNAKE_CASE__								=   ()

											SCREAMING_SNAKE_CASE__								=   self.output(self.after_norm(__A		)		)
											SCREAMING_SNAKE_CASE__								=   self.layer_scale_a.unsqueeze(-1		).unsqueeze(-1		) * layer_output
											# Second residual connection
											SCREAMING_SNAKE_CASE__								=   hidden_states + self.drop_path(__A		)

											SCREAMING_SNAKE_CASE__								=   (output,) + outputs
											return outputs

									else:
											SCREAMING_SNAKE_CASE__								=   self.drop_path(self.pooling(self.before_norm(__A		)		)		)
											# First residual connection
											SCREAMING_SNAKE_CASE__								=   pooling_output + hidden_states
											SCREAMING_SNAKE_CASE__								=   ()

											# Second residual connection inside the PoolFormerOutput block
											SCREAMING_SNAKE_CASE__								=   self.drop_path(self.output(self.after_norm(__A		)		)		)
											SCREAMING_SNAKE_CASE__								=   hidden_states + layer_output

											SCREAMING_SNAKE_CASE__								=   (output,) + outputs
											return outputs




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Union[str, Any]				,					__A     :List[Any]		) -> Union[str, Any]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   config
									# stochastic depth decay rule
									SCREAMING_SNAKE_CASE__								=   [x.item() for x in torch.linspace(0				,					config.drop_path_rate				,					sum(config.depths		)		)]

									# patch embeddings
									SCREAMING_SNAKE_CASE__								=   []
									for i in range(config.num_encoder_blocks		):
											embeddings.append(
											    PoolFormerEmbeddings(
											        patch_size=config.patch_sizes[i]				,					stride=config.strides[i]				,					padding=config.padding[i]				,					num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]				,					hidden_size=config.hidden_sizes[i]				,					)		)
									SCREAMING_SNAKE_CASE__								=   nn.ModuleList(__A		)

									# Transformer blocks
									SCREAMING_SNAKE_CASE__								=   []
									SCREAMING_SNAKE_CASE__								=   0
									for i in range(config.num_encoder_blocks		):
											# each block consists of layers
											SCREAMING_SNAKE_CASE__								=   []
											if i != 0:
													cur += config.depths[i - 1]
											for j in range(config.depths[i]		):
													layers.append(
													    PoolFormerLayer(
													        __A				,					num_channels=config.hidden_sizes[i]				,					pool_size=config.pool_size				,					hidden_size=config.hidden_sizes[i]				,					intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio		)				,					drop_path=dpr[cur + j]				,					)		)
											blocks.append(nn.ModuleList(__A		)		)

									SCREAMING_SNAKE_CASE__								=   nn.ModuleList(__A		)


							def   _snake_case	(				self     :str				,					__A     :Tuple				,					__A     :Dict=False				,					__A     :Tuple=True		) -> Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   () if output_hidden_states else None

									SCREAMING_SNAKE_CASE__								=   pixel_values
									for idx, layers in enumerate(zip(self.patch_embeddings				,					self.block		)		):
											SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   layers
											# Get patch embeddings from hidden_states
											SCREAMING_SNAKE_CASE__								=   embedding_layer(__A		)
											# Send the embeddings through the blocks
											for _, blk in enumerate(__A		):
													SCREAMING_SNAKE_CASE__								=   blk(__A		)
													SCREAMING_SNAKE_CASE__								=   layer_outputs[0]

											if output_hidden_states:
													SCREAMING_SNAKE_CASE__								=   all_hidden_states + (hidden_states,)

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

									return BaseModelOutputWithNoAttention(last_hidden_state=__A				,					hidden_states=__A		)




class 	UpperCamelCase_							(  UpperCamelCase__     ):
							lowerCamelCase_							     =       PoolFormerConfig
							lowerCamelCase_							     =       "poolformer"
							lowerCamelCase_							     =       "pixel_values"
							lowerCamelCase_							     =       True


							def   _snake_case	(				self     :Optional[Any]				,					__A     :Tuple		) -> Dict:



									"""simple docstring"""



									if isinstance(__A				,					(nn.Linear, nn.Convad)		):
											module.weight.data.normal_(mean=0.0				,					std=self.config.initializer_range		)
											if module.bias is not None:
													module.bias.data.zero_()
									elif isinstance(__A				,					nn.LayerNorm		):
											module.bias.data.zero_()
											module.weight.data.fill_(1.0		)


							def   _snake_case	(				self     :str				,					__A     :Optional[Any]				,					__A     :Union[str, Any]=False		) -> Any:



									"""simple docstring"""



									if isinstance(__A				,					__A		):
											SCREAMING_SNAKE_CASE__								=   value



_lowerCamelCase 					=      R'\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it 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 ([`PoolFormerConfig`]): 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'

_lowerCamelCase 					=      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            [`PoolFormerImageProcessor.__call__`] for details.\n'




@add_start_docstrings(
    "The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top."					,	UpperCamelCase__					,	)
class 	UpperCamelCase_							(  UpperCamelCase__     ):


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



									"""simple docstring"""



									super().__init__(__A		)
									SCREAMING_SNAKE_CASE__								=   config

									SCREAMING_SNAKE_CASE__								=   PoolFormerEncoder(__A		)

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


							def   _snake_case	(				self     :Union[str, Any]		) -> Union[str, Any]:



									"""simple docstring"""



									return self.embeddings.patch_embeddings


							@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   _snake_case	(				self     :Dict				,					__A     :Optional[torch.FloatTensor] = None				,					__A     :Optional[bool] = None				,					__A     :Optional[bool] = None				,					) -> Union[Tuple, BaseModelOutputWithNoAttention]:



									"""simple docstring"""



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

									if pixel_values is None:
											raise ValueError("""You have to specify pixel_values"""		)

									SCREAMING_SNAKE_CASE__								=   self.encoder(
									    __A				,					output_hidden_states=__A				,					return_dict=__A				,					)
									SCREAMING_SNAKE_CASE__								=   encoder_outputs[0]

									if not return_dict:
											return (sequence_output, None) + encoder_outputs[1:]

									return BaseModelOutputWithNoAttention(
									    last_hidden_state=__A				,					hidden_states=encoder_outputs.hidden_states				,					)




class 	UpperCamelCase_							(  nn.Module     ):


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



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   nn.Linear(config.hidden_size				,					config.hidden_size		)


							def   _snake_case	(				self     :List[Any]				,					__A     :Dict		) -> int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.dense(__A		)
									return output




@add_start_docstrings(
    "\n    PoolFormer Model transformer with an image classification head on top\n    "					,	UpperCamelCase__					,	)
class 	UpperCamelCase_							(  UpperCamelCase__     ):


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



									"""simple docstring"""



									super().__init__(__A		)
									SCREAMING_SNAKE_CASE__								=   config.num_labels
									SCREAMING_SNAKE_CASE__								=   PoolFormerModel(__A		)

									# Final norm
									SCREAMING_SNAKE_CASE__								=   PoolFormerGroupNorm(config.hidden_sizes[-1]		)
									# Classifier head
									SCREAMING_SNAKE_CASE__								=   (
									    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   _snake_case	(				self     :int				,					__A     :Optional[torch.FloatTensor] = None				,					__A     :Optional[torch.LongTensor] = None				,					__A     :Optional[bool] = None				,					__A     :Optional[bool] = None				,					) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:



									"""simple docstring"""



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

									SCREAMING_SNAKE_CASE__								=   self.poolformer(
									    __A				,					output_hidden_states=__A				,					return_dict=__A				,					)

									SCREAMING_SNAKE_CASE__								=   outputs[0]

									SCREAMING_SNAKE_CASE__								=   self.classifier(self.norm(__A		).mean([-2, -1]		)		)

									SCREAMING_SNAKE_CASE__								=   None
									if labels is not None:
											if self.config.problem_type is None:
													if self.num_labels == 1:
															SCREAMING_SNAKE_CASE__								=   """regression"""
													elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
															SCREAMING_SNAKE_CASE__								=   """single_label_classification"""
													else:
															SCREAMING_SNAKE_CASE__								=   """multi_label_classification"""

											if self.config.problem_type == "regression":
													SCREAMING_SNAKE_CASE__								=   MSELoss()
													if self.num_labels == 1:
															SCREAMING_SNAKE_CASE__								=   loss_fct(logits.squeeze()				,					labels.squeeze()		)
													else:
															SCREAMING_SNAKE_CASE__								=   loss_fct(__A				,					__A		)
											elif self.config.problem_type == "single_label_classification":
													SCREAMING_SNAKE_CASE__								=   CrossEntropyLoss()
													SCREAMING_SNAKE_CASE__								=   loss_fct(logits.view(-1				,					self.num_labels		)				,					labels.view(-1		)		)
											elif self.config.problem_type == "multi_label_classification":
													SCREAMING_SNAKE_CASE__								=   BCEWithLogitsLoss()
													SCREAMING_SNAKE_CASE__								=   loss_fct(__A				,					__A		)

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

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







'''simple docstring'''







import collections.abc
from typing import Optional, Tuple, Union

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

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


__lowerCAmelCase        =							logging.get_logger(__name__)

# General docstring
__lowerCAmelCase        =							"PoolFormerConfig"

# Base docstring
__lowerCAmelCase        =							"sail/poolformer_s12"
__lowerCAmelCase        =							[1, 5_12, 7, 7]

# Image classification docstring
__lowerCAmelCase        =							"sail/poolformer_s12"
__lowerCAmelCase        =							"tabby, tabby cat"

__lowerCAmelCase        =							[
    "sail/poolformer_s12",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]


def       __UpperCamelCase					(      lowercase_       :  List[Any]   ,       lowercase_       :  float = 0.0   ,       lowercase_       :  bool = False   ):
       """simple docstring"""




       if drop_prob == 0.0 or not training:
              return input
       a_						     =			1 - drop_prob
       a_						     =			(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
       a_						     =			keep_prob + torch.rand(lowercase_   ,       dtype=input.dtype   ,       device=input.device   )
       random_tensor.floor_()  # binarize
       a_						     =			input.div(lowercase_   ) * random_tensor
       return output




class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__ = None	):






         """simple docstring"""


         super().__init__()
         a_						     =			drop_prob







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


         return drop_path(UpperCamelCase__   ,	self.drop_prob   ,	self.training	)







  def    _a       (			self	):






         """simple docstring"""


         return "p={}".format(self.drop_prob	)





class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__=None	):






         """simple docstring"""


         super().__init__()
         a_						     =			patch_size if isinstance(UpperCamelCase__   ,	collections.abc.Iterable	) else (patch_size, patch_size)
         a_						     =			stride if isinstance(UpperCamelCase__   ,	collections.abc.Iterable	) else (stride, stride)
         a_						     =			padding if isinstance(UpperCamelCase__   ,	collections.abc.Iterable	) else (padding, padding)

         a_						     =			nn.Convad(UpperCamelCase__   ,	UpperCamelCase__   ,	kernel_size=UpperCamelCase__   ,	stride=UpperCamelCase__   ,	padding=UpperCamelCase__	)
         a_						     =			norm_layer(UpperCamelCase__	) if norm_layer else nn.Identity()







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


         a_						     =			self.projection(UpperCamelCase__	)
         a_						     =			self.norm(UpperCamelCase__	)
         return embeddings





class   __SCREAMING_SNAKE_CASE       (nn.GroupNorm       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__   ,	**UpperCamelCase__	):






         """simple docstring"""


         super().__init__(1   ,	UpperCamelCase__   ,	**UpperCamelCase__	)





class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__	):






         """simple docstring"""


         super().__init__()
         a_						     =			nn.AvgPoolad(UpperCamelCase__   ,	stride=1   ,	padding=pool_size // 2   ,	count_include_pad=UpperCamelCase__	)







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


         return self.pool(UpperCamelCase__	) - hidden_states





class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__	):






         """simple docstring"""


         super().__init__()
         a_						     =			nn.Convad(UpperCamelCase__   ,	UpperCamelCase__   ,	1	)
         a_						     =			nn.Convad(UpperCamelCase__   ,	UpperCamelCase__   ,	1	)
         a_						     =			PoolFormerDropPath(UpperCamelCase__	)
         if isinstance(config.hidden_act   ,	UpperCamelCase__	):
                a_						     =			ACTaFN[config.hidden_act]
         else:
                a_						     =			config.hidden_act







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


         a_						     =			self.conva(UpperCamelCase__	)
         a_						     =			self.act_fn(UpperCamelCase__	)
         a_						     =			self.drop(UpperCamelCase__	)
         a_						     =			self.conva(UpperCamelCase__	)
         a_						     =			self.drop(UpperCamelCase__	)

         return hidden_states





class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__	):






         """simple docstring"""


         super().__init__()
         a_						     =			PoolFormerPooling(UpperCamelCase__	)
         a_						     =			PoolFormerOutput(UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__	)
         a_						     =			PoolFormerGroupNorm(UpperCamelCase__	)
         a_						     =			PoolFormerGroupNorm(UpperCamelCase__	)

         # Useful for training neural nets
         a_						     =			PoolFormerDropPath(UpperCamelCase__	) if drop_path > 0.0 else nn.Identity()
         a_						     =			config.use_layer_scale
         if config.use_layer_scale:
                a_						     =			nn.Parameter(
                    config.layer_scale_init_value * torch.ones((UpperCamelCase__)	)   ,	requires_grad=UpperCamelCase__	)
                a_						     =			nn.Parameter(
                    config.layer_scale_init_value * torch.ones((UpperCamelCase__)	)   ,	requires_grad=UpperCamelCase__	)







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


         if self.use_layer_scale:
                a_						     =			self.pooling(self.before_norm(UpperCamelCase__	)	)
                a_						     =			self.layer_scale_a.unsqueeze(-1	).unsqueeze(-1	) * pooling_output
                # First residual connection
                a_						     =			hidden_states + self.drop_path(UpperCamelCase__	)
                a_						     =			()

                a_						     =			self.output(self.after_norm(UpperCamelCase__	)	)
                a_						     =			self.layer_scale_a.unsqueeze(-1	).unsqueeze(-1	) * layer_output
                # Second residual connection
                a_						     =			hidden_states + self.drop_path(UpperCamelCase__	)

                a_						     =			(output,) + outputs
                return outputs

         else:
                a_						     =			self.drop_path(self.pooling(self.before_norm(UpperCamelCase__	)	)	)
                # First residual connection
                a_						     =			pooling_output + hidden_states
                a_						     =			()

                # Second residual connection inside the PoolFormerOutput block
                a_						     =			self.drop_path(self.output(self.after_norm(UpperCamelCase__	)	)	)
                a_						     =			hidden_states + layer_output

                a_						     =			(output,) + outputs
                return outputs





class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__	):






         """simple docstring"""


         super().__init__()
         a_						     =			config
         # stochastic depth decay rule
         a_						     =			[x.item() for x in torch.linspace(0   ,	config.drop_path_rate   ,	sum(config.depths	)	)]

         # patch embeddings
         a_						     =			[]
         for i in range(config.num_encoder_blocks	):
                embeddings.append(
                    PoolFormerEmbeddings(
                        patch_size=config.patch_sizes[i]   ,	stride=config.strides[i]   ,	padding=config.padding[i]   ,	num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]   ,	hidden_size=config.hidden_sizes[i]   ,	)	)
         a_						     =			nn.ModuleList(UpperCamelCase__	)

         # Transformer blocks
         a_						     =			[]
         a_						     =			0
         for i in range(config.num_encoder_blocks	):
                # each block consists of layers
                a_						     =			[]
                if i != 0:
                       cur += config.depths[i - 1]
                for j in range(config.depths[i]	):
                       layers.append(
                           PoolFormerLayer(
                               UpperCamelCase__   ,	num_channels=config.hidden_sizes[i]   ,	pool_size=config.pool_size   ,	hidden_size=config.hidden_sizes[i]   ,	intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio	)   ,	drop_path=dpr[cur + j]   ,	)	)
                blocks.append(nn.ModuleList(UpperCamelCase__	)	)

         a_						     =			nn.ModuleList(UpperCamelCase__	)







  def    _a       (			self   ,	UpperCamelCase__   ,	UpperCamelCase__=False   ,	UpperCamelCase__=True	):






         """simple docstring"""


         a_						     =			() if output_hidden_states else None

         a_						     =			pixel_values
         for idx, layers in enumerate(zip(self.patch_embeddings   ,	self.block	)	):
                a_     ,	a_						     =			layers
                # Get patch embeddings from hidden_states
                a_						     =			embedding_layer(UpperCamelCase__	)
                # Send the embeddings through the blocks
                for _, blk in enumerate(UpperCamelCase__	):
                       a_						     =			blk(UpperCamelCase__	)
                       a_						     =			layer_outputs[0]

                if output_hidden_states:
                       a_						     =			all_hidden_states + (hidden_states,)

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

         return BaseModelOutputWithNoAttention(last_hidden_state=UpperCamelCase__   ,	hidden_states=UpperCamelCase__	)





class   __SCREAMING_SNAKE_CASE       (__A       ):


  """simple docstring"""


  _a  : int      			=     PoolFormerConfig
  _a  : Optional[Any]      			=     '''poolformer'''
  _a  : Union[str, Any]      			=     '''pixel_values'''
  _a  : Optional[int]      			=     True







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


         if isinstance(UpperCamelCase__   ,	(nn.Linear, nn.Convad)	):
                module.weight.data.normal_(mean=0.0   ,	std=self.config.initializer_range	)
                if module.bias is not None:
                       module.bias.data.zero_()
         elif isinstance(UpperCamelCase__   ,	nn.LayerNorm	):
                module.bias.data.zero_()
                module.weight.data.fill_(1.0	)







  def    _a       (			self   ,	UpperCamelCase__   ,	UpperCamelCase__=False	):






         """simple docstring"""


         if isinstance(UpperCamelCase__   ,	UpperCamelCase__	):
                a_						     =			value





__lowerCAmelCase        =							R"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it 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 ([`PoolFormerConfig`]): 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"

__lowerCAmelCase        =							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            [`PoolFormerImageProcessor.__call__`] for details.\n"




@add_start_docstrings(
    '''The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top.'''							,       __A							,       )
class   __SCREAMING_SNAKE_CASE       (__A       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__	):






         """simple docstring"""


         super().__init__(UpperCamelCase__	)
         a_						     =			config

         a_						     =			PoolFormerEncoder(UpperCamelCase__	)

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







  def    _a       (			self	):






         """simple docstring"""


         return self.embeddings.patch_embeddings







  @add_start_docstrings_to_model_forward(UpperCamelCase__	)
  @add_code_sample_docstrings(
      checkpoint=_CHECKPOINT_FOR_DOC   ,	output_type=UpperCamelCase__   ,	config_class=_CONFIG_FOR_DOC   ,	modality='vision'   ,	expected_output=_EXPECTED_OUTPUT_SHAPE   ,	)
  def    _a       (			self   ,	UpperCamelCase__ = None   ,	UpperCamelCase__ = None   ,	UpperCamelCase__ = None   ,	):






         """simple docstring"""


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

         if pixel_values is None:
                raise ValueError('You have to specify pixel_values'	)

         a_						     =			self.encoder(
             UpperCamelCase__   ,	output_hidden_states=UpperCamelCase__   ,	return_dict=UpperCamelCase__   ,	)
         a_						     =			encoder_outputs[0]

         if not return_dict:
                return (sequence_output, None) + encoder_outputs[1:]

         return BaseModelOutputWithNoAttention(
             last_hidden_state=UpperCamelCase__   ,	hidden_states=encoder_outputs.hidden_states   ,	)





class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__	):






         """simple docstring"""


         super().__init__()
         a_						     =			nn.Linear(config.hidden_size   ,	config.hidden_size	)







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


         a_						     =			self.dense(UpperCamelCase__	)
         return output





@add_start_docstrings(
    '''
    PoolFormer Model transformer with an image classification head on top
    '''							,       __A							,       )
class   __SCREAMING_SNAKE_CASE       (__A       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__	):






         """simple docstring"""


         super().__init__(UpperCamelCase__	)
         a_						     =			config.num_labels
         a_						     =			PoolFormerModel(UpperCamelCase__	)

         # Final norm
         a_						     =			PoolFormerGroupNorm(config.hidden_sizes[-1]	)
         # Classifier head
         a_						     =			(
             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(UpperCamelCase__	)
  @add_code_sample_docstrings(
      checkpoint=_IMAGE_CLASS_CHECKPOINT   ,	output_type=UpperCamelCase__   ,	config_class=_CONFIG_FOR_DOC   ,	expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,	)
  def    _a       (			self   ,	UpperCamelCase__ = None   ,	UpperCamelCase__ = None   ,	UpperCamelCase__ = None   ,	UpperCamelCase__ = None   ,	):






         """simple docstring"""


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

         a_						     =			self.poolformer(
             UpperCamelCase__   ,	output_hidden_states=UpperCamelCase__   ,	return_dict=UpperCamelCase__   ,	)

         a_						     =			outputs[0]

         a_						     =			self.classifier(self.norm(UpperCamelCase__	).mean([-2, -1]	)	)

         a_						     =			None
         if labels is not None:
                if self.config.problem_type is None:
                       if self.num_labels == 1:
                              a_						     =			'regression'
                       elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                              a_						     =			'single_label_classification'
                       else:
                              a_						     =			'multi_label_classification'

                if self.config.problem_type == "regression":
                       a_						     =			MSELoss()
                       if self.num_labels == 1:
                              a_						     =			loss_fct(logits.squeeze()   ,	labels.squeeze()	)
                       else:
                              a_						     =			loss_fct(UpperCamelCase__   ,	UpperCamelCase__	)
                elif self.config.problem_type == "single_label_classification":
                       a_						     =			CrossEntropyLoss()
                       a_						     =			loss_fct(logits.view(-1   ,	self.num_labels	)   ,	labels.view(-1	)	)
                elif self.config.problem_type == "multi_label_classification":
                       a_						     =			BCEWithLogitsLoss()
                       a_						     =			loss_fct(UpperCamelCase__   ,	UpperCamelCase__	)

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

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




