




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

# General docstring
A     ='PoolFormerConfig'

# Base docstring
A     ='sail/poolformer_s12'
A     =[1, 5_12, 7, 7]

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

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





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


class 			_a  (	nn.Module ):
       def __init__(	self						:	str   ,				lowercase						:	Optional[float] = None				):



           '''simple docstring'''

           super().__init__()
           UpperCAmelCase			       =		drop_prob
       def     A      (	self						:	Union[str, Any]   ,				lowercase						:	torch.Tensor				):



           '''simple docstring'''

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







       def     A      (	self						:	Optional[Any]				):



           '''simple docstring'''

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






class 			_a  (	nn.Module ):
       def __init__(	self						:	Optional[int]   ,				lowercase						:	Optional[int]   ,				lowercase						:	int   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	Optional[int]   ,				lowercase						:	int   ,				lowercase						:	Tuple=None				):



           '''simple docstring'''

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

           UpperCAmelCase			       =		nn.Convad(lowercase   ,				lowercase   ,				kernel_size=lowercase   ,				stride=lowercase   ,				padding=lowercase				)
           UpperCAmelCase			       =		norm_layer(lowercase				) if norm_layer else nn.Identity()







       def     A      (	self						:	int   ,				lowercase						:	Optional[int]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.projection(lowercase				)
           UpperCAmelCase			       =		self.norm(lowercase				)
           return embeddings






class 			_a  (	nn.GroupNorm ):
       def __init__(	self						:	Optional[int]   ,				lowercase						:	Optional[Any]   ,				**lowercase						:	int				):



           '''simple docstring'''

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






class 			_a  (	nn.Module ):
       def __init__(	self						:	Dict   ,				lowercase						:	Dict				):



           '''simple docstring'''

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







       def     A      (	self						:	str   ,				lowercase						:	Dict				):



           '''simple docstring'''

           return self.pool(lowercase				) - hidden_states






class 			_a  (	nn.Module ):
       def __init__(	self						:	Any   ,				lowercase						:	List[str]   ,				lowercase						:	str   ,				lowercase						:	Optional[int]   ,				lowercase						:	Any				):



           '''simple docstring'''

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







       def     A      (	self						:	List[str]   ,				lowercase						:	int				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.conva(lowercase				)
           UpperCAmelCase			       =		self.act_fn(lowercase				)
           UpperCAmelCase			       =		self.drop(lowercase				)
           UpperCAmelCase			       =		self.conva(lowercase				)
           UpperCAmelCase			       =		self.drop(lowercase				)

           return hidden_states






class 			_a  (	nn.Module ):
       def __init__(	self						:	Tuple   ,				lowercase						:	Optional[int]   ,				lowercase						:	List[Any]   ,				lowercase						:	Dict   ,				lowercase						:	List[Any]   ,				lowercase						:	Any   ,				lowercase						:	List[Any]				):



           '''simple docstring'''

           super().__init__()
           UpperCAmelCase			       =		PoolFormerPooling(lowercase				)
           UpperCAmelCase			       =		PoolFormerOutput(lowercase   ,				lowercase   ,				lowercase   ,				lowercase				)
           UpperCAmelCase			       =		PoolFormerGroupNorm(lowercase				)
           UpperCAmelCase			       =		PoolFormerGroupNorm(lowercase				)

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







       def     A      (	self						:	Tuple   ,				lowercase						:	str				):



           '''simple docstring'''

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

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

               UpperCAmelCase			       =		(output,) + outputs
               return outputs

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

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

               UpperCAmelCase			       =		(output,) + outputs
               return outputs






class 			_a  (	nn.Module ):
       def __init__(	self						:	Optional[int]   ,				lowercase						:	Tuple				):



           '''simple docstring'''

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

           # patch embeddings
           UpperCAmelCase			       =		[]
           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]   ,				)				)
           UpperCAmelCase			       =		nn.ModuleList(lowercase				)

           # Transformer blocks
           UpperCAmelCase			       =		[]
           UpperCAmelCase			       =		0
           for i in range(config.num_encoder_blocks				):
               # each block consists of layers
               UpperCAmelCase			       =		[]
               if i != 0:
                   cur += config.depths[i - 1]
               for j in range(config.depths[i]				):
                   layers.append(
                       PoolFormerLayer(
                           lowercase   ,				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(lowercase				)				)

           UpperCAmelCase			       =		nn.ModuleList(lowercase				)







       def     A      (	self						:	Optional[Any]   ,				lowercase						:	Optional[int]   ,				lowercase						:	Union[str, Any]=False   ,				lowercase						:	List[Any]=True				):



           '''simple docstring'''

           UpperCAmelCase			       =		() if output_hidden_states else None

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

               if output_hidden_states:
                   UpperCAmelCase			       =		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=lowercase   ,				hidden_states=lowercase				)






class 			_a  (	__a ):
       __a				:     Optional[int]   				=					PoolFormerConfig
       __a				:     Dict   				=					"""poolformer"""
       __a				:     Union[str, Any]   				=					"""pixel_values"""
       __a				:     Dict   				=					True
       def     A      (	self						:	Tuple   ,				lowercase						:	Any				):



           '''simple docstring'''

           if isinstance(lowercase   ,				(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(lowercase   ,				nn.LayerNorm				):
               module.bias.data.zero_()
               module.weight.data.fill_(1.0				)







       def     A      (	self						:	str   ,				lowercase						:	str   ,				lowercase						:	Dict=False				):



           '''simple docstring'''

           if isinstance(lowercase   ,				lowercase				):
               UpperCAmelCase			       =		value






A     =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'

A     =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 			_a  (	__a ):
       def __init__(	self						:	str   ,				lowercase						:	Optional[Any]				):



           '''simple docstring'''

           super().__init__(lowercase				)
           UpperCAmelCase			       =		config

           UpperCAmelCase			       =		PoolFormerEncoder(lowercase				)

           # Initialize weights and apply final processing
           self.post_init()
       def     A      (	self						:	int				):



           '''simple docstring'''

           return self.embeddings.patch_embeddings







       @add_start_docstrings_to_model_forward(lowercase				)
       @add_code_sample_docstrings(
           checkpoint=_CHECKPOINT_FOR_DOC   ,				output_type=lowercase   ,				config_class=_CONFIG_FOR_DOC   ,				modality='''vision'''   ,				expected_output=_EXPECTED_OUTPUT_SHAPE   ,				)
       def     A      (	self						:	Tuple   ,				lowercase						:	Optional[torch.FloatTensor] = None   ,				lowercase						:	Optional[bool] = None   ,				lowercase						:	Optional[bool] = None   ,				):



           '''simple docstring'''

           UpperCAmelCase			       =		(
               output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
           )
           UpperCAmelCase			       =		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'''				)

           UpperCAmelCase			       =		self.encoder(
               lowercase   ,				output_hidden_states=lowercase   ,				return_dict=lowercase   ,				)
           UpperCAmelCase			       =		encoder_outputs[0]

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

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






class 			_a  (	nn.Module ):
       def __init__(	self						:	Optional[int]   ,				lowercase						:	str				):



           '''simple docstring'''

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







       def     A      (	self						:	str   ,				lowercase						:	List[str]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.dense(lowercase				)
           return output






@add_start_docstrings(
    """
    PoolFormer Model transformer with an image classification head on top
    """							,							__a							,							)
class 			_a  (	__a ):
       def __init__(	self						:	str   ,				lowercase						:	List[Any]				):



           '''simple docstring'''

           super().__init__(lowercase				)
           UpperCAmelCase			       =		config.num_labels
           UpperCAmelCase			       =		PoolFormerModel(lowercase				)

           # Final norm
           UpperCAmelCase			       =		PoolFormerGroupNorm(config.hidden_sizes[-1]				)
           # Classifier head
           UpperCAmelCase			       =		(
               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(lowercase				)
       @add_code_sample_docstrings(
           checkpoint=_IMAGE_CLASS_CHECKPOINT   ,				output_type=lowercase   ,				config_class=_CONFIG_FOR_DOC   ,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,				)
       def     A      (	self						:	List[Any]   ,				lowercase						:	Optional[torch.FloatTensor] = None   ,				lowercase						:	Optional[torch.LongTensor] = None   ,				lowercase						:	Optional[bool] = None   ,				lowercase						:	Optional[bool] = None   ,				):



           '''simple docstring'''

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

           UpperCAmelCase			       =		self.poolformer(
               lowercase   ,				output_hidden_states=lowercase   ,				return_dict=lowercase   ,				)

           UpperCAmelCase			       =		outputs[0]

           UpperCAmelCase			       =		self.classifier(self.norm(lowercase				).mean([-2, -1]				)				)

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

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

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

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











"""simple docstring"""



from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


a_ =			logging.get_logger(__name__)

a_ =			{
    """shi-labs/nat-mini-in1k-224""": """https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json""",
    # See all Nat models at https://huggingface.co/models?filter=nat
}






class 		__snake_case					(		SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 _lowerCamelCase        =     """nat"""

 _lowerCamelCase        =     {
     """num_attention_heads""": """num_heads""",
     """num_hidden_layers""": """num_layers""",
 }

 def __init__(    self					,						__lowerCamelCase=4					,						__lowerCamelCase=3					,						__lowerCamelCase=64					,						__lowerCamelCase=[3, 4, 6, 5]					,						__lowerCamelCase=[2, 4, 8, 16]					,						__lowerCamelCase=7					,						__lowerCamelCase=3.0					,						__lowerCamelCase=True					,						__lowerCamelCase=0.0					,						__lowerCamelCase=0.0					,						__lowerCamelCase=0.1					,						__lowerCamelCase="gelu"					,						__lowerCamelCase=0.0_2					,						__lowerCamelCase=1e-5					,						__lowerCamelCase=0.0					,						__lowerCamelCase=None					,						__lowerCamelCase=None					,						**__lowerCamelCase					,						):



     '''simple docstring'''




     super().__init__(**__lowerCamelCase  )

     __A    :						Union[str, Any]						 =     patch_size
     __A    :						Optional[Any]						 =     num_channels
     __A    :						Tuple						 =     embed_dim
     __A    :						Dict						 =     depths
     __A    :						str						 =     len(__lowerCamelCase  )
     __A    :						Optional[Any]						 =     num_heads
     __A    :						str						 =     kernel_size
     __A    :						Any						 =     mlp_ratio
     __A    :						Optional[int]						 =     qkv_bias
     __A    :						str						 =     hidden_dropout_prob
     __A    :						Any						 =     attention_probs_dropout_prob
     __A    :						int						 =     drop_path_rate
     __A    :						int						 =     hidden_act
     __A    :						Any						 =     layer_norm_eps
     __A    :						Tuple						 =     initializer_range
     # we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
     # this indicates the channel dimension after the last stage of the model
     __A    :						int						 =     int(embed_dim * 2 ** (len(__lowerCamelCase  ) - 1)  )
     __A    :						Union[str, Any]						 =     layer_scale_init_value
     __A    :						List[str]						 =     ['''stem'''] + [F"""stage{idx}""" for idx in range(1					,						len(__lowerCamelCase  ) + 1  )]
     __A ,       __A    :						Any						 =     get_aligned_output_features_output_indices(
         out_features=__lowerCamelCase					,						out_indices=__lowerCamelCase					,						stage_names=self.stage_names  )




