

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

# General docstring
_lowerCAmelCase			:						int  					=     'PoolFormerConfig'

# Base docstring
_lowerCAmelCase			:						str  					=     'sail/poolformer_s12'
_lowerCAmelCase			:						Any  					=     [1, 5_1_2, 7, 7]

# Image classification docstring
_lowerCAmelCase			:						List[Any]  					=     'sail/poolformer_s12'
_lowerCAmelCase			:						str  					=     'tabby, tabby cat'

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







def a_		(							UpperCamelCase_ :				Optional[int] ,       UpperCamelCase_ :				float = 0.0 ,       UpperCamelCase_ :				bool = False	)				-> Optional[int]:







 """simple docstring"""

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


class      lowerCAmelCase					(     nn.Module					):


       '''simple docstring'''







       def __init__(						self       :	int    ,						__snake_case       :	Optional[float] = None  ) ->			List[str]:







        '''simple docstring'''

        super().__init__()
        lowerCamelCase       =     drop_prob

       def 	lowerCamelCase__    (						self       :	int    ,						__snake_case       :	torch.Tensor  ) ->			Optional[Any]:







        '''simple docstring'''

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





       def 	lowerCamelCase__    (						self       :	List[Any]  ) ->			Optional[int]:







        '''simple docstring'''

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


class      lowerCAmelCase					(     nn.Module					):


       '''simple docstring'''







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







        '''simple docstring'''

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

        lowerCamelCase       =     nn.Convad(__lowerCamelCase    ,						__lowerCamelCase    ,						kernel_size=__lowerCamelCase    ,						stride=__lowerCamelCase    ,						padding=__lowerCamelCase  )
        lowerCamelCase       =     norm_layer(__lowerCamelCase  ) if norm_layer else nn.Identity()





       def 	lowerCamelCase__    (						self       :	Union[str, Any]    ,						__snake_case       :	Any  ) ->			Tuple:







        '''simple docstring'''

        lowerCamelCase       =     self.projection(__lowerCamelCase  )
        lowerCamelCase       =     self.norm(__lowerCamelCase  )
        return embeddings


class      lowerCAmelCase					(     nn.GroupNorm					):


       '''simple docstring'''







       def __init__(						self       :	Dict    ,						__snake_case       :	Any    ,						**__snake_case       :	List[str]  ) ->			Optional[Any]:







        '''simple docstring'''

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


class      lowerCAmelCase					(     nn.Module					):


       '''simple docstring'''







       def __init__(						self       :	Optional[Any]    ,						__snake_case       :	Dict  ) ->			Union[str, Any]:







        '''simple docstring'''

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





       def 	lowerCamelCase__    (						self       :	int    ,						__snake_case       :	int  ) ->			List[Any]:







        '''simple docstring'''

        return self.pool(__lowerCamelCase  ) - hidden_states


class      lowerCAmelCase					(     nn.Module					):


       '''simple docstring'''







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







        '''simple docstring'''

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





       def 	lowerCamelCase__    (						self       :	Any    ,						__snake_case       :	int  ) ->			Optional[int]:







        '''simple docstring'''

        lowerCamelCase       =     self.conva(__lowerCamelCase  )
        lowerCamelCase       =     self.act_fn(__lowerCamelCase  )
        lowerCamelCase       =     self.drop(__lowerCamelCase  )
        lowerCamelCase       =     self.conva(__lowerCamelCase  )
        lowerCamelCase       =     self.drop(__lowerCamelCase  )

        return hidden_states


class      lowerCAmelCase					(     nn.Module					):


       '''simple docstring'''







       def __init__(						self       :	str    ,						__snake_case       :	int    ,						__snake_case       :	List[str]    ,						__snake_case       :	int    ,						__snake_case       :	int    ,						__snake_case       :	Dict    ,						__snake_case       :	List[str]  ) ->			Optional[int]:







        '''simple docstring'''

        super().__init__()
        lowerCamelCase       =     PoolFormerPooling(__lowerCamelCase  )
        lowerCamelCase       =     PoolFormerOutput(__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase  )
        lowerCamelCase       =     PoolFormerGroupNorm(__lowerCamelCase  )
        lowerCamelCase       =     PoolFormerGroupNorm(__lowerCamelCase  )

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





       def 	lowerCamelCase__    (						self       :	Optional[int]    ,						__snake_case       :	Optional[int]  ) ->			Tuple:







        '''simple docstring'''

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

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

         lowerCamelCase       =     (output,) + outputs
         return outputs

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

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

         lowerCamelCase       =     (output,) + outputs
         return outputs


class      lowerCAmelCase					(     nn.Module					):


       '''simple docstring'''







       def __init__(						self       :	Any    ,						__snake_case       :	Dict  ) ->			Union[str, Any]:







        '''simple docstring'''

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

        # patch embeddings
        lowerCamelCase       =     []
        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       =     nn.ModuleList(__lowerCamelCase  )

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

        lowerCamelCase       =     nn.ModuleList(__lowerCamelCase  )





       def 	lowerCamelCase__    (						self       :	Optional[int]    ,						__snake_case       :	Optional[Any]    ,						__snake_case       :	Union[str, Any]=False    ,						__snake_case       :	List[Any]=True  ) ->			Dict:







        '''simple docstring'''

        lowerCamelCase       =     () if output_hidden_states else None

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

         if output_hidden_states:
          lowerCamelCase       =     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=__lowerCamelCase    ,						hidden_states=__lowerCamelCase  )


class      lowerCAmelCase					(     __snake_case					):


       '''simple docstring'''







       snake_case     				=						PoolFormerConfig
       snake_case     				=						'poolformer'
       snake_case     				=						'pixel_values'
       snake_case     				=						True

       def 	lowerCamelCase__    (						self       :	Optional[Any]    ,						__snake_case       :	Any  ) ->			Optional[Any]:







        '''simple docstring'''

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





       def 	lowerCamelCase__    (						self       :	str    ,						__snake_case       :	Tuple    ,						__snake_case       :	Any=False  ) ->			Optional[int]:







        '''simple docstring'''

        if isinstance(__lowerCamelCase    ,						__lowerCamelCase  ):
         lowerCamelCase       =     value

_lowerCAmelCase			:						Optional[int]  					=     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			:						Dict  					=     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.' ,					__snake_case ,					)
class      lowerCAmelCase					(     __snake_case					):


       '''simple docstring'''







       def __init__(						self       :	str    ,						__snake_case       :	str  ) ->			Optional[Any]:







        '''simple docstring'''

        super().__init__(__lowerCamelCase  )
        lowerCamelCase       =     config

        lowerCamelCase       =     PoolFormerEncoder(__lowerCamelCase  )

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

       def 	lowerCamelCase__    (						self       :	Any  ) ->			List[str]:







        '''simple docstring'''

        return self.embeddings.patch_embeddings





       @add_start_docstrings_to_model_forward(__lowerCamelCase  )
       @add_code_sample_docstrings(
           checkpoint=_CHECKPOINT_FOR_DOC    ,						output_type=__lowerCamelCase    ,						config_class=_CONFIG_FOR_DOC    ,						modality='vision'    ,						expected_output=_EXPECTED_OUTPUT_SHAPE    ,						)
       def 	lowerCamelCase__    (						self       :	Union[str, Any]    ,						__snake_case       :	Optional[torch.FloatTensor] = None    ,						__snake_case       :	Optional[bool] = None    ,						__snake_case       :	Optional[bool] = None    ,						) ->			List[str]:







        '''simple docstring'''

        lowerCamelCase       =     (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        lowerCamelCase       =     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       =     self.encoder(
            __lowerCamelCase    ,						output_hidden_states=__lowerCamelCase    ,						return_dict=__lowerCamelCase    ,						)
        lowerCamelCase       =     encoder_outputs[0]

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

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


class      lowerCAmelCase					(     nn.Module					):


       '''simple docstring'''







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







        '''simple docstring'''

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





       def 	lowerCamelCase__    (						self       :	Tuple    ,						__snake_case       :	Tuple  ) ->			Tuple:







        '''simple docstring'''

        lowerCamelCase       =     self.dense(__lowerCamelCase  )
        return output


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


       '''simple docstring'''







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







        '''simple docstring'''

        super().__init__(__lowerCamelCase  )
        lowerCamelCase       =     config.num_labels
        lowerCamelCase       =     PoolFormerModel(__lowerCamelCase  )

        # Final norm
        lowerCamelCase       =     PoolFormerGroupNorm(config.hidden_sizes[-1]  )
        # Classifier head
        lowerCamelCase       =     (
            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(__lowerCamelCase  )
       @add_code_sample_docstrings(
           checkpoint=_IMAGE_CLASS_CHECKPOINT    ,						output_type=__lowerCamelCase    ,						config_class=_CONFIG_FOR_DOC    ,						expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT    ,						)
       def 	lowerCamelCase__    (						self       :	Union[str, Any]    ,						__snake_case       :	Optional[torch.FloatTensor] = None    ,						__snake_case       :	Optional[torch.LongTensor] = None    ,						__snake_case       :	Optional[bool] = None    ,						__snake_case       :	Optional[bool] = None    ,						) ->			Dict:







        '''simple docstring'''

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

        lowerCamelCase       =     self.poolformer(
            __lowerCamelCase    ,						output_hidden_states=__lowerCamelCase    ,						return_dict=__lowerCamelCase    ,						)

        lowerCamelCase       =     outputs[0]

        lowerCamelCase       =     self.classifier(self.norm(__lowerCamelCase  ).mean([-2, -1]  )  )

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

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

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

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







from typing import Any

import numpy as np





def 	__a     (      A__			:  np.ndarray       ):
    return np.array_equal(A__			,  matrix.conjugate().T       )





def 	__a     (      A__			:  np.ndarray			,  A__			:  np.ndarray       ):
    SCREAMING_SNAKE_CASE										=   v.conjugate().T
    SCREAMING_SNAKE_CASE										=   v_star.dot(A__       )
    assert isinstance(A__			,  np.ndarray       )
    return (v_star_dot.dot(A__       )) / (v_star.dot(A__       ))





def 	__a     (      ):
    SCREAMING_SNAKE_CASE										=   np.array([[2, 2 + 1j, 4], [2 - 1j, 3, 1j], [4, -1j, 1]]       )
    SCREAMING_SNAKE_CASE										=   np.array([[1], [2], [3]]       )
    assert is_hermitian(A__       ), F"{a} is not hermitian."
    print(rayleigh_quotient(A__			,  A__       )       )

    SCREAMING_SNAKE_CASE										=   np.array([[1, 2, 4], [2, 3, -1], [4, -1, 1]]       )
    assert is_hermitian(A__       ), F"{a} is not hermitian."
    assert rayleigh_quotient(A__			,  A__       ) == float(3       )


if __name__ == "__main__":
    import doctest

    doctest.testmod()
    tests()