





import math
import sys



def     _lowerCAmelCase		(  _lowerCAmelCase     ):
			'''simple docstring'''


			if number != int(_lowerCAmelCase     ):
						raise ValueError("""the value of input must be a natural number"""     )
			if number < 0:
						raise ValueError("""the value of input must not be a negative number"""     )
			if number == 0:
						return 1
			A_     :				int			      =	[-1] * (number + 1)
			A_     :				List[Any]			      =	0
			for i in range(1  ,number + 1     ):
						A_     :				Tuple			      =	sys.maxsize
						A_     :				int			      =	int(math.sqrt(_lowerCAmelCase     )     )
						for j in range(1  ,root + 1     ):
									A_     :				Any			      =	1 + answers[i - (j**2)]
									A_     :				Tuple			      =	min(_lowerCAmelCase  ,_lowerCAmelCase     )
						A_     :				Union[str, Any]			      =	answer
			return answers[number]


if __name__ == "__main__":
						import doctest

						doctest.testmod()













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     _lowerCAmelCase		(  _lowerCAmelCase  ,_lowerCAmelCase = 0.0  ,_lowerCAmelCase = False     ):
   '''simple docstring'''


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

class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__ = None       ):
         super().__init__()
         A_     :				int			      =	drop_prob
      def 		_lowerCamelCase      (		self      ,       a__       ):
         return drop_path(a__      ,       self.drop_prob      ,       self.training       )






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






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__=None       ):
         super().__init__()
         A_     :				Dict			      =	patch_size if isinstance(a__      ,       collections.abc.Iterable       ) else (patch_size, patch_size)
         A_     :				Optional[Any]			      =	stride if isinstance(a__      ,       collections.abc.Iterable       ) else (stride, stride)
         A_     :				int			      =	padding if isinstance(a__      ,       collections.abc.Iterable       ) else (padding, padding)

         A_     :				Dict			      =	nn.Convad(a__      ,       a__      ,       kernel_size=a__      ,       stride=a__      ,       padding=a__       )
         A_     :				Optional[Any]			      =	norm_layer(a__       ) if norm_layer else nn.Identity()






      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				Dict			      =	self.projection(a__       )
         A_     :				List[Any]			      =	self.norm(a__       )
         return embeddings






class    _UpperCAmelCase							(		nn.GroupNorm     ):
      def __init__(		self      ,       a__      ,       **a__       ):
         super().__init__(1      ,       a__      ,       **a__       )






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__       ):
         super().__init__()
         A_     :				Optional[int]			      =	nn.AvgPoolad(a__      ,       stride=1      ,       padding=pool_size // 2      ,       count_include_pad=a__       )






      def 		_lowerCamelCase      (		self      ,       a__       ):
         return self.pool(a__       ) - hidden_states






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__      ,       a__      ,       a__      ,       a__       ):
         super().__init__()
         A_     :				str			      =	nn.Convad(a__      ,       a__      ,       1       )
         A_     :				List[str]			      =	nn.Convad(a__      ,       a__      ,       1       )
         A_     :				Optional[Any]			      =	PoolFormerDropPath(a__       )
         if isinstance(config.hidden_act      ,       a__       ):
            A_     :				int			      =	ACTaFN[config.hidden_act]
         else:
            A_     :				Optional[int]			      =	config.hidden_act






      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				Tuple			      =	self.conva(a__       )
         A_     :				str			      =	self.act_fn(a__       )
         A_     :				Union[str, Any]			      =	self.drop(a__       )
         A_     :				int			      =	self.conva(a__       )
         A_     :				Any			      =	self.drop(a__       )

         return hidden_states






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__       ):
         super().__init__()
         A_     :				Dict			      =	PoolFormerPooling(a__       )
         A_     :				Dict			      =	PoolFormerOutput(a__      ,       a__      ,       a__      ,       a__       )
         A_     :				Dict			      =	PoolFormerGroupNorm(a__       )
         A_     :				Dict			      =	PoolFormerGroupNorm(a__       )

         # Useful for training neural nets
         A_     :				Union[str, Any]			      =	PoolFormerDropPath(a__       ) if drop_path > 0.0 else nn.Identity()
         A_     :				Union[str, Any]			      =	config.use_layer_scale
         if config.use_layer_scale:
            A_     :				Tuple			      =	nn.Parameter(
                config.layer_scale_init_value * torch.ones((a__)       )      ,       requires_grad=a__       )
            A_     :				Tuple			      =	nn.Parameter(
                config.layer_scale_init_value * torch.ones((a__)       )      ,       requires_grad=a__       )






      def 		_lowerCamelCase      (		self      ,       a__       ):
         if self.use_layer_scale:
            A_     :				List[str]			      =	self.pooling(self.before_norm(a__       )       )
            A_     :				Optional[Any]			      =	self.layer_scale_a.unsqueeze(-1       ).unsqueeze(-1       ) * pooling_output
            # First residual connection
            A_     :				List[Any]			      =	hidden_states + self.drop_path(a__       )
            A_     :				Tuple			      =	()

            A_     :				List[Any]			      =	self.output(self.after_norm(a__       )       )
            A_     :				str			      =	self.layer_scale_a.unsqueeze(-1       ).unsqueeze(-1       ) * layer_output
            # Second residual connection
            A_     :				Optional[Any]			      =	hidden_states + self.drop_path(a__       )

            A_     :				int			      =	(output,) + outputs
            return outputs

         else:
            A_     :				Union[str, Any]			      =	self.drop_path(self.pooling(self.before_norm(a__       )       )       )
            # First residual connection
            A_     :				Any			      =	pooling_output + hidden_states
            A_     :				Tuple			      =	()

            # Second residual connection inside the PoolFormerOutput block
            A_     :				Dict			      =	self.drop_path(self.output(self.after_norm(a__       )       )       )
            A_     :				int			      =	hidden_states + layer_output

            A_     :				str			      =	(output,) + outputs
            return outputs






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__       ):
         super().__init__()
         A_     :				Union[str, Any]			      =	config
         # stochastic depth decay rule
         A_     :				Union[str, Any]			      =	[x.item() for x in torch.linspace(0      ,       config.drop_path_rate      ,       sum(config.depths       )       )]

         # patch embeddings
         A_     :				Union[str, Any]			      =	[]
         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_     :				Optional[int]			      =	nn.ModuleList(a__       )

         # Transformer blocks
         A_     :				Union[str, Any]			      =	[]
         A_     :				Union[str, Any]			      =	0
         for i in range(config.num_encoder_blocks       ):
            # each block consists of layers
            A_     :				Dict			      =	[]
            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__       )       )

         A_     :				str			      =	nn.ModuleList(a__       )






      def 		_lowerCamelCase      (		self      ,       a__      ,       a__=False      ,       a__=True       ):
         A_     :				Optional[int]			      =	() if output_hidden_states else None

         A_     :				str			      =	pixel_values
         for idx, layers in enumerate(zip(self.patch_embeddings      ,       self.block       )       ):
            A_      ,				A_     :				Union[str, Any]			      =	layers
            # Get patch embeddings from hidden_states
            A_     :				Tuple			      =	embedding_layer(a__       )
            # Send the embeddings through the blocks
            for _, blk in enumerate(a__       ):
               A_     :				List[str]			      =	blk(a__       )
               A_     :				Union[str, Any]			      =	layer_outputs[0]

            if output_hidden_states:
               A_     :				Dict			      =	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							(		_lowerCamelCase     ):
      a						       =							PoolFormerConfig
      a						       =							'''poolformer'''
      a						       =							'''pixel_values'''
      a						       =							True
      def 		_lowerCamelCase      (		self      ,       a__       ):
         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 		_lowerCamelCase      (		self      ,       a__      ,       a__=False       ):
         if isinstance(a__      ,       a__       ):
            A_     :				List[Any]			      =	value






_lowerCAmelCase 				=     R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`PoolFormerConfig`]): 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.
"""

_lowerCAmelCase 				=     R"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`PoolFormerImageProcessor.__call__`] for details.
"""

@add_start_docstrings(
    '''The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top.'''			,       _lowerCamelCase			,       )
class    _UpperCAmelCase							(		_lowerCamelCase     ):
      def __init__(		self      ,       a__       ):
         super().__init__(a__       )
         A_     :				List[Any]			      =	config

         A_     :				Optional[int]			      =	PoolFormerEncoder(a__       )

         # Initialize weights and apply final processing
         self.post_init()
      def 		_lowerCamelCase      (		self       ):
         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 		_lowerCamelCase      (		self      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       ):
         A_     :				Union[str, Any]			      =	(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         A_     :				str			      =	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_     :				Optional[Any]			      =	self.encoder(
             a__      ,       output_hidden_states=a__      ,       return_dict=a__      ,       )
         A_     :				Optional[int]			      =	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      ,       a__       ):
         super().__init__()
         A_     :				int			      =	nn.Linear(config.hidden_size      ,       config.hidden_size       )






      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				Any			      =	self.dense(a__       )
         return output






@add_start_docstrings(
    '''
    PoolFormer Model transformer with an image classification head on top
    '''			,       _lowerCamelCase			,       )
class    _UpperCAmelCase							(		_lowerCamelCase     ):
      def __init__(		self      ,       a__       ):
         super().__init__(a__       )
         A_     :				str			      =	config.num_labels
         A_     :				Union[str, Any]			      =	PoolFormerModel(a__       )

         # Final norm
         A_     :				str			      =	PoolFormerGroupNorm(config.hidden_sizes[-1]       )
         # Classifier head
         A_     :				Optional[int]			      =	(
             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 		_lowerCamelCase      (		self      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       ):
         A_     :				List[Any]			      =	return_dict if return_dict is not None else self.config.use_return_dict

         A_     :				Optional[int]			      =	self.poolformer(
             a__      ,       output_hidden_states=a__      ,       return_dict=a__      ,       )

         A_     :				List[str]			      =	outputs[0]

         A_     :				int			      =	self.classifier(self.norm(a__       ).mean([-2, -1]       )       )

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

            if self.config.problem_type == "regression":
               A_     :				int			      =	MSELoss()
               if self.num_labels == 1:
                  A_     :				List[str]			      =	loss_fct(logits.squeeze()      ,       labels.squeeze()       )
               else:
                  A_     :				Tuple			      =	loss_fct(a__      ,       a__       )
            elif self.config.problem_type == "single_label_classification":
               A_     :				str			      =	CrossEntropyLoss()
               A_     :				int			      =	loss_fct(logits.view(-1      ,       self.num_labels       )      ,       labels.view(-1       )       )
            elif self.config.problem_type == "multi_label_classification":
               A_     :				List[str]			      =	BCEWithLogitsLoss()
               A_     :				Optional[Any]			      =	loss_fct(a__      ,       a__       )

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

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






