




"""simple docstring"""
def 						UpperCAmelCase				(			A__:    List[str]							,				A__:    Optional[Any]     )  ->     List[Any]:
			__lowerCamelCase					:	Any	       =				0
			while b > 0:
						if b & 1:
									res += a

						a += a
						b >>= 1

			return res





def 						UpperCAmelCase				(			A__:    Any							,				A__:    Optional[int]							,				A__:    Union[str, Any]     )  ->     Optional[int]:
			__lowerCamelCase					:	Optional[int]	       =				0
			while b > 0:
						if b & 1:
									__lowerCamelCase					:	Dict	       =				((res % c) + (a % c)) % c

						a += a
						b >>= 1

			return res









"""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


a_				:    Tuple										=     logging.get_logger(__name__)

# General docstring
a_				:    List[str]										=     '''PoolFormerConfig'''

# Base docstring
a_				:    Optional[Any]										=     '''sail/poolformer_s12'''
a_				:    List[Any]										=     [1, 5_12, 7, 7]

# Image classification docstring
a_				:    Any										=     '''sail/poolformer_s12'''
a_				:    Optional[int]										=     '''tabby, tabby cat'''

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





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

class __lowercase(      nn.Module				):




					'''simple docstring'''


					def __init__(    self						,	__a = None     ):
								super().__init__()
								__lowerCamelCase					:	int	       =				drop_prob
					def     snake_case_				(    self						,	__a     ):
								return drop_path(__a						,	self.drop_prob						,	self.training     )





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



class __lowercase(      nn.Module				):




					'''simple docstring'''


					def __init__(    self						,	__a						,	__a						,	__a						,	__a						,	__a						,	__a=None     ):
								super().__init__()
								__lowerCamelCase					:	int	       =				patch_size if isinstance(__a						,	collections.abc.Iterable     ) else (patch_size, patch_size)
								__lowerCamelCase					:	int	       =				stride if isinstance(__a						,	collections.abc.Iterable     ) else (stride, stride)
								__lowerCamelCase					:	Optional[int]	       =				padding if isinstance(__a						,	collections.abc.Iterable     ) else (padding, padding)

								__lowerCamelCase					:	Optional[Any]	       =				nn.Convad(__a						,	__a						,	kernel_size=__a						,	stride=__a						,	padding=__a     )
								__lowerCamelCase					:	List[str]	       =				norm_layer(__a     ) if norm_layer else nn.Identity()





					def     snake_case_				(    self						,	__a     ):
								__lowerCamelCase					:	List[Any]	       =				self.projection(__a     )
								__lowerCamelCase					:	Dict	       =				self.norm(__a     )
								return embeddings



class __lowercase(      nn.GroupNorm				):




					'''simple docstring'''


					def __init__(    self						,	__a						,	**__a     ):
								super().__init__(1						,	__a						,	**__a     )



class __lowercase(      nn.Module				):




					'''simple docstring'''


					def __init__(    self						,	__a     ):
								super().__init__()
								__lowerCamelCase					:	str	       =				nn.AvgPoolad(__a						,	stride=1						,	padding=pool_size // 2						,	count_include_pad=__a     )





					def     snake_case_				(    self						,	__a     ):
								return self.pool(__a     ) - hidden_states



class __lowercase(      nn.Module				):




					'''simple docstring'''


					def __init__(    self						,	__a						,	__a						,	__a						,	__a     ):
								super().__init__()
								__lowerCamelCase					:	Any	       =				nn.Convad(__a						,	__a						,	1     )
								__lowerCamelCase					:	Dict	       =				nn.Convad(__a						,	__a						,	1     )
								__lowerCamelCase					:	List[Any]	       =				PoolFormerDropPath(__a     )
								if isinstance(config.hidden_act						,	__a     ):
											__lowerCamelCase					:	List[str]	       =				ACTaFN[config.hidden_act]
								else:
											__lowerCamelCase					:	str	       =				config.hidden_act





					def     snake_case_				(    self						,	__a     ):
								__lowerCamelCase					:	int	       =				self.conva(__a     )
								__lowerCamelCase					:	Dict	       =				self.act_fn(__a     )
								__lowerCamelCase					:	List[str]	       =				self.drop(__a     )
								__lowerCamelCase					:	int	       =				self.conva(__a     )
								__lowerCamelCase					:	str	       =				self.drop(__a     )

								return hidden_states



class __lowercase(      nn.Module				):




					'''simple docstring'''


					def __init__(    self						,	__a						,	__a						,	__a						,	__a						,	__a						,	__a     ):
								super().__init__()
								__lowerCamelCase					:	Tuple	       =				PoolFormerPooling(__a     )
								__lowerCamelCase					:	Union[str, Any]	       =				PoolFormerOutput(__a						,	__a						,	__a						,	__a     )
								__lowerCamelCase					:	List[Any]	       =				PoolFormerGroupNorm(__a     )
								__lowerCamelCase					:	List[Any]	       =				PoolFormerGroupNorm(__a     )

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





					def     snake_case_				(    self						,	__a     ):
								if self.use_layer_scale:
											__lowerCamelCase					:	Union[str, Any]	       =				self.pooling(self.before_norm(__a     )     )
											__lowerCamelCase					:	Any	       =				self.layer_scale_a.unsqueeze(-1     ).unsqueeze(-1     ) * pooling_output
											# First residual connection
											__lowerCamelCase					:	Optional[Any]	       =				hidden_states + self.drop_path(__a     )
											__lowerCamelCase					:	Tuple	       =				()

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

											__lowerCamelCase					:	Optional[Any]	       =				(output,) + outputs
											return outputs

								else:
											__lowerCamelCase					:	Tuple	       =				self.drop_path(self.pooling(self.before_norm(__a     )     )     )
											# First residual connection
											__lowerCamelCase					:	List[str]	       =				pooling_output + hidden_states
											__lowerCamelCase					:	int	       =				()

											# Second residual connection inside the PoolFormerOutput block
											__lowerCamelCase					:	List[str]	       =				self.drop_path(self.output(self.after_norm(__a     )     )     )
											__lowerCamelCase					:	str	       =				hidden_states + layer_output

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



class __lowercase(      nn.Module				):




					'''simple docstring'''


					def __init__(    self						,	__a     ):
								super().__init__()
								__lowerCamelCase					:	int	       =				config
								# stochastic depth decay rule
								__lowerCamelCase					:	int	       =				[x.item() for x in torch.linspace(0						,	config.drop_path_rate						,	sum(config.depths     )     )]

								# patch embeddings
								__lowerCamelCase					:	List[str]	       =				[]
								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]						,	)     )
								__lowerCamelCase					:	Optional[int]	       =				nn.ModuleList(__a     )

								# Transformer blocks
								__lowerCamelCase					:	Any	       =				[]
								__lowerCamelCase					:	int	       =				0
								for i in range(config.num_encoder_blocks     ):
											# each block consists of layers
											__lowerCamelCase					:	Optional[int]	       =				[]
											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     )     )

								__lowerCamelCase					:	str	       =				nn.ModuleList(__a     )





					def     snake_case_				(    self						,	__a						,	__a=False						,	__a=True     ):
								__lowerCamelCase					:	Union[str, Any]	       =				() if output_hidden_states else None

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

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




					'''simple docstring'''


					__a						:						Tuple     		= PoolFormerConfig
					__a						:						Tuple     		= 'poolformer'
					__a						:						Optional[int]     		= 'pixel_values'
					__a						:						Optional[Any]     		= True
					def     snake_case_				(    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     snake_case_				(    self						,	__a						,	__a=False     ):
								if isinstance(__a						,	__a     ):
											__lowerCamelCase					:	Union[str, Any]	       =				value



a_				:    Union[str, Any]										=     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.
'''

a_				:    List[str]										=     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.'	,	lowercase__	,	)
class __lowercase(      lowercase__				):




					'''simple docstring'''


					def __init__(    self						,	__a     ):
								super().__init__(__a     )
								__lowerCamelCase					:	Optional[Any]	       =				config

								__lowerCamelCase					:	Any	       =				PoolFormerEncoder(__a     )

								# Initialize weights and apply final processing
								self.post_init()
					def     snake_case_				(    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     snake_case_				(    self						,	__a = None						,	__a = None						,	__a = None						,	):
								__lowerCamelCase					:	Union[str, Any]	       =				(
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								__lowerCamelCase					:	Dict	       =				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'     )

								__lowerCamelCase					:	Any	       =				self.encoder(
								    __a						,	output_hidden_states=__a						,	return_dict=__a						,	)
								__lowerCamelCase					:	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 __lowercase(      nn.Module				):




					'''simple docstring'''


					def __init__(    self						,	__a     ):
								super().__init__()
								__lowerCamelCase					:	Optional[Any]	       =				nn.Linear(config.hidden_size						,	config.hidden_size     )





					def     snake_case_				(    self						,	__a     ):
								__lowerCamelCase					:	List[Any]	       =				self.dense(__a     )
								return output



@add_start_docstrings(
    '\n    PoolFormer Model transformer with an image classification head on top\n    '	,	lowercase__	,	)
class __lowercase(      lowercase__				):




					'''simple docstring'''


					def __init__(    self						,	__a     ):
								super().__init__(__a     )
								__lowerCamelCase					:	str	       =				config.num_labels
								__lowerCamelCase					:	Optional[Any]	       =				PoolFormerModel(__a     )

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

								__lowerCamelCase					:	Tuple	       =				self.poolformer(
								    __a						,	output_hidden_states=__a						,	return_dict=__a						,	)

								__lowerCamelCase					:	int	       =				outputs[0]

								__lowerCamelCase					:	Optional[int]	       =				self.classifier(self.norm(__a     ).mean([-2, -1]     )     )

								__lowerCamelCase					:	Union[str, Any]	       =				None
								if labels is not None:
											if self.config.problem_type is None:
														if self.num_labels == 1:
																	__lowerCamelCase					:	Any	       =				'regression'
														elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																	__lowerCamelCase					:	Any	       =				'single_label_classification'
														else:
																	__lowerCamelCase					:	Optional[Any]	       =				'multi_label_classification'

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

								if not return_dict:
											__lowerCamelCase					:	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     )



