



"""simple docstring"""






import inspect
import unittest

from transformers import SegformerConfig, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         MODEL_MAPPING,
         SegformerForImageClassification,
         SegformerForSemanticSegmentation,
         SegformerModel,
     )
     from transformers.models.segformer.modeling_segformer import SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import SegformerImageProcessor



class     UpperCAmelCase    (	__SCREAMING_SNAKE_CASE  ):



   def    __UpperCAmelCase    (       self    :	List[Any]				):



     """simple docstring"""
     _snake_case   	=      self.config_class(**self.inputs_dict				)
     self.parent.assertTrue(hasattr(__lowerCamelCase							,       '''hidden_sizes'''				)				)
     self.parent.assertTrue(hasattr(__lowerCamelCase							,       '''num_attention_heads'''				)				)
     self.parent.assertTrue(hasattr(__lowerCamelCase							,       '''num_encoder_blocks'''				)				)



class     UpperCAmelCase    :



   def __init__(       self    :	str							,       __lowerCamelCase    :	Union[str, Any]							,       __lowerCamelCase    :	Optional[int]=1_3							,       __lowerCamelCase    :	Optional[int]=6_4							,       __lowerCamelCase    :	str=3							,       __lowerCamelCase    :	Optional[int]=4							,       __lowerCamelCase    :	List[str]=[2, 2, 2, 2]							,       __lowerCamelCase    :	List[str]=[8, 4, 2, 1]							,       __lowerCamelCase    :	List[str]=[1_6, 3_2, 6_4, 1_2_8]							,       __lowerCamelCase    :	int=[1, 4, 8, 1_6]							,       __lowerCamelCase    :	Dict=[1, 2, 4, 8]							,       __lowerCamelCase    :	Union[str, Any]=True							,       __lowerCamelCase    :	int=True							,       __lowerCamelCase    :	Tuple="gelu"							,       __lowerCamelCase    :	Any=0.1							,       __lowerCamelCase    :	str=0.1							,       __lowerCamelCase    :	int=0.0_2							,       __lowerCamelCase    :	Any=3							,       __lowerCamelCase    :	str=None							,       ):



     """simple docstring"""
     _snake_case   	=      parent
     _snake_case   	=      batch_size
     _snake_case   	=      image_size
     _snake_case   	=      num_channels
     _snake_case   	=      num_encoder_blocks
     _snake_case   	=      sr_ratios
     _snake_case   	=      depths
     _snake_case   	=      hidden_sizes
     _snake_case   	=      downsampling_rates
     _snake_case   	=      num_attention_heads
     _snake_case   	=      is_training
     _snake_case   	=      use_labels
     _snake_case   	=      hidden_act
     _snake_case   	=      hidden_dropout_prob
     _snake_case   	=      attention_probs_dropout_prob
     _snake_case   	=      initializer_range
     _snake_case   	=      num_labels
     _snake_case   	=      scope



   def    __UpperCAmelCase    (       self    :	List[Any]				):



     """simple docstring"""
     _snake_case   	=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

     _snake_case   	=      None
     if self.use_labels:
       _snake_case   	=      ids_tensor([self.batch_size, self.image_size, self.image_size]							,       self.num_labels				)

     _snake_case   	=      self.get_config()
     return config, pixel_values, labels



   def    __UpperCAmelCase    (       self    :	int				):



     """simple docstring"""
     return SegformerConfig(
         image_size=self.image_size							,       num_channels=self.num_channels							,       num_encoder_blocks=self.num_encoder_blocks							,       depths=self.depths							,       hidden_sizes=self.hidden_sizes							,       num_attention_heads=self.num_attention_heads							,       hidden_act=self.hidden_act							,       hidden_dropout_prob=self.hidden_dropout_prob							,       attention_probs_dropout_prob=self.attention_probs_dropout_prob							,       initializer_range=self.initializer_range							,       )



   def    __UpperCAmelCase    (       self    :	Optional[Any]							,       __lowerCamelCase    :	Tuple							,       __lowerCamelCase    :	int							,       __lowerCamelCase    :	List[str]				):



     """simple docstring"""
     _snake_case   	=      SegformerModel(config=__lowerCamelCase				)
     model.to(__lowerCamelCase				)
     model.eval()
     _snake_case   	=      model(__lowerCamelCase				)
     _snake_case   	=      _snake_case   	=      self.image_size // (self.downsampling_rates[-1] * 2)
     self.parent.assertEqual(
         result.last_hidden_state.shape							,       (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)				)



   def    __UpperCAmelCase    (       self    :	Any							,       __lowerCamelCase    :	Dict							,       __lowerCamelCase    :	Any							,       __lowerCamelCase    :	List[str]				):



     """simple docstring"""
     _snake_case   	=      self.num_labels
     _snake_case   	=      SegformerForSemanticSegmentation(__lowerCamelCase				)
     model.to(__lowerCamelCase				)
     model.eval()
     _snake_case   	=      model(__lowerCamelCase				)
     self.parent.assertEqual(
         result.logits.shape							,       (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)				)
     _snake_case   	=      model(__lowerCamelCase							,       labels=__lowerCamelCase				)
     self.parent.assertEqual(
         result.logits.shape							,       (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)				)
     self.parent.assertGreater(result.loss							,       0.0				)



   def    __UpperCAmelCase    (       self    :	Optional[Any]							,       __lowerCamelCase    :	Union[str, Any]							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	str				):



     """simple docstring"""
     _snake_case   	=      1
     _snake_case   	=      SegformerForSemanticSegmentation(config=__lowerCamelCase				)
     model.to(__lowerCamelCase				)
     model.eval()
     _snake_case   	=      torch.randint(0							,       1							,       (self.batch_size, self.image_size, self.image_size)				).to(__lowerCamelCase				)
     _snake_case   	=      model(__lowerCamelCase							,       labels=__lowerCamelCase				)
     self.parent.assertGreater(result.loss							,       0.0				)





   def    __UpperCAmelCase    (       self    :	Union[str, Any]				):



     """simple docstring"""
     _snake_case   	=      self.prepare_config_and_inputs()
     _snake_case	,     _snake_case	,     _snake_case   	=      config_and_inputs
     _snake_case   	=      {'''pixel_values''': pixel_values}
     return config, inputs_dict



@require_torch
class     UpperCAmelCase    (	__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,unittest.TestCase  ):
   A__		:	str           =  (
       (
           SegformerModel,
           SegformerForSemanticSegmentation,
           SegformerForImageClassification,
       )
       if is_torch_available()
       else ()
   )
   A__		:	Optional[int]           =  (
       {
           '''feature-extraction''': SegformerModel,
           '''image-classification''': SegformerForImageClassification,
           '''image-segmentation''': SegformerForSemanticSegmentation,
       }
       if is_torch_available()
       else {}
   )

   A__		:	Optional[Any]           =  True
   A__		:	List[str]           =  False
   A__		:	Any           =  False
   A__		:	str           =  False



   def    __UpperCAmelCase    (       self    :	Optional[int]				):



     """simple docstring"""
     _snake_case   	=      SegformerModelTester(self				)
     _snake_case   	=      SegformerConfigTester(self							,       config_class=__lowerCamelCase				)



   def    __UpperCAmelCase    (       self    :	List[str]				):



     """simple docstring"""
     self.config_tester.run_common_tests()



   def    __UpperCAmelCase    (       self    :	Optional[int]				):



     """simple docstring"""
     _snake_case   	=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__lowerCamelCase				)



   def    __UpperCAmelCase    (       self    :	Tuple				):



     """simple docstring"""
     _snake_case   	=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_binary_image_segmentation(*__lowerCamelCase				)



   def    __UpperCAmelCase    (       self    :	Optional[Any]				):



     """simple docstring"""
     _snake_case   	=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_segmentation(*__lowerCamelCase				)



   @unittest.skip('''SegFormer does not use inputs_embeds'''				)
   def    __UpperCAmelCase    (       self    :	Any				):



     """simple docstring"""
     pass



   @unittest.skip('''SegFormer does not have get_input_embeddings method and get_output_embeddings methods'''				)
   def    __UpperCAmelCase    (       self    :	List[str]				):



     """simple docstring"""
     pass



   def    __UpperCAmelCase    (       self    :	Optional[Any]				):



     """simple docstring"""
     _snake_case	,     _snake_case   	=      self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       _snake_case   	=      model_class(__lowerCamelCase				)
       _snake_case   	=      inspect.signature(model.forward				)
       # signature.parameters is an OrderedDict => so arg_names order is deterministic
       _snake_case   	=      [*signature.parameters.keys()]

       _snake_case   	=      ['''pixel_values''']
       self.assertListEqual(arg_names[:1]							,       __lowerCamelCase				)



   def    __UpperCAmelCase    (       self    :	Optional[Any]				):



     """simple docstring"""
     _snake_case	,     _snake_case   	=      self.model_tester.prepare_config_and_inputs_for_common()
     _snake_case   	=      True

     for model_class in self.all_model_classes:
       _snake_case   	=      True
       _snake_case   	=      False
       _snake_case   	=      True
       _snake_case   	=      model_class(__lowerCamelCase				)
       model.to(__lowerCamelCase				)
       model.eval()
       with torch.no_grad():
         _snake_case   	=      model(**self._prepare_for_class(__lowerCamelCase							,       __lowerCamelCase				)				)
       _snake_case   	=      outputs.attentions

       _snake_case   	=      sum(self.model_tester.depths				)
       self.assertEqual(len(__lowerCamelCase				)							,       __lowerCamelCase				)

       # check that output_attentions also work using config
       del inputs_dict["output_attentions"]
       _snake_case   	=      True
       _snake_case   	=      model_class(__lowerCamelCase				)
       model.to(__lowerCamelCase				)
       model.eval()
       with torch.no_grad():
         _snake_case   	=      model(**self._prepare_for_class(__lowerCamelCase							,       __lowerCamelCase				)				)
       _snake_case   	=      outputs.attentions

       self.assertEqual(len(__lowerCamelCase				)							,       __lowerCamelCase				)

       # verify the first attentions (first block, first layer)
       _snake_case   	=      (self.model_tester.image_size // 4) ** 2
       _snake_case   	=      (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
       self.assertListEqual(
           list(attentions[0].shape[-3:]				)							,       [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]							,       )

       # verify the last attentions (last block, last layer)
       _snake_case   	=      (self.model_tester.image_size // 3_2) ** 2
       _snake_case   	=      (self.model_tester.image_size // (3_2 * self.model_tester.sr_ratios[-1])) ** 2
       self.assertListEqual(
           list(attentions[-1].shape[-3:]				)							,       [self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]							,       )
       _snake_case   	=      len(__lowerCamelCase				)

       # Check attention is always last and order is fine
       _snake_case   	=      True
       _snake_case   	=      True
       _snake_case   	=      model_class(__lowerCamelCase				)
       model.to(__lowerCamelCase				)
       model.eval()
       with torch.no_grad():
         _snake_case   	=      model(**self._prepare_for_class(__lowerCamelCase							,       __lowerCamelCase				)				)

       self.assertEqual(out_len + 1							,       len(__lowerCamelCase				)				)

       _snake_case   	=      outputs.attentions

       self.assertEqual(len(__lowerCamelCase				)							,       __lowerCamelCase				)
       # verify the first attentions (first block, first layer)
       _snake_case   	=      (self.model_tester.image_size // 4) ** 2
       _snake_case   	=      (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
       self.assertListEqual(
           list(self_attentions[0].shape[-3:]				)							,       [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]							,       )



   def    __UpperCAmelCase    (       self    :	Optional[Any]				):



     """simple docstring"""
     def check_hidden_states_output(__lowerCamelCase    :	Optional[Any]							,       __lowerCamelCase    :	List[Any]							,       __lowerCamelCase    :	List[str]				):
       _snake_case   	=      model_class(__lowerCamelCase				)
       model.to(__lowerCamelCase				)
       model.eval()

       with torch.no_grad():
         _snake_case   	=      model(**self._prepare_for_class(__lowerCamelCase							,       __lowerCamelCase				)				)

       _snake_case   	=      outputs.hidden_states

       _snake_case   	=      self.model_tester.num_encoder_blocks
       self.assertEqual(len(__lowerCamelCase				)							,       __lowerCamelCase				)

       # verify the first hidden states (first block)
       self.assertListEqual(
           list(hidden_states[0].shape[-3:]				)							,       [
               self.model_tester.hidden_sizes[0],
               self.model_tester.image_size // 4,
               self.model_tester.image_size // 4,
           ]							,       )

     _snake_case	,     _snake_case   	=      self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       _snake_case   	=      True
       check_hidden_states_output(__lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase				)

       # check that output_hidden_states also work using config
       del inputs_dict["output_hidden_states"]
       _snake_case   	=      True

       check_hidden_states_output(__lowerCamelCase							,       __lowerCamelCase							,       __lowerCamelCase				)



   def    __UpperCAmelCase    (       self    :	Dict				):



     """simple docstring"""
     if not self.model_tester.is_training:
       return

     _snake_case	,     _snake_case   	=      self.model_tester.prepare_config_and_inputs_for_common()
     _snake_case   	=      True

     for model_class in self.all_model_classes:
       if model_class in get_values(__lowerCamelCase				):
         continue

       _snake_case   	=      model_class(__lowerCamelCase				)
       model.to(__lowerCamelCase				)
       model.train()
       _snake_case   	=      self._prepare_for_class(__lowerCamelCase							,       __lowerCamelCase							,       return_labels=__lowerCamelCase				)
       _snake_case   	=      model(**__lowerCamelCase				).loss
       loss.backward()



   @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''				)
   def    __UpperCAmelCase    (       self    :	List[Any]				):



     """simple docstring"""
     pass





   @slow
   def    __UpperCAmelCase    (       self    :	Optional[Any]				):



     """simple docstring"""
     for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _snake_case   	=      SegformerModel.from_pretrained(__lowerCamelCase				)
       self.assertIsNotNone(__lowerCamelCase				)




def 			snake_case		(		)						->				int:
  _snake_case   	=      Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''	)
  return image







@require_torch
class     UpperCAmelCase    (	unittest.TestCase  ):



   @slow
   def    __UpperCAmelCase    (       self    :	int				):



     """simple docstring"""
     # only resize + normalize
     _snake_case   	=      SegformerImageProcessor(
         image_scale=(5_1_2, 5_1_2)							,       keep_ratio=__lowerCamelCase							,       align=__lowerCamelCase							,       do_random_crop=__lowerCamelCase				)
     _snake_case   	=      SegformerForSemanticSegmentation.from_pretrained('''nvidia/segformer-b0-finetuned-ade-512-512'''				).to(
         __lowerCamelCase				)

     _snake_case   	=      prepare_img()
     _snake_case   	=      image_processor(images=__lowerCamelCase							,       return_tensors='''pt'''				)
     _snake_case   	=      encoded_inputs.pixel_values.to(__lowerCamelCase				)

     with torch.no_grad():
       _snake_case   	=      model(__lowerCamelCase				)

     _snake_case   	=      torch.Size((1, model.config.num_labels, 1_2_8, 1_2_8)				)
     self.assertEqual(outputs.logits.shape							,       __lowerCamelCase				)

     _snake_case   	=      torch.tensor(
         [
             [[-4.6_3_1_0, -5.5_2_3_2, -6.2_3_5_6], [-5.1_9_2_1, -6.1_4_4_4, -6.5_9_9_6], [-5.4_4_2_4, -6.2_7_9_0, -6.7_5_7_4]],
             [[-1_2.1_3_9_1, -1_3.3_1_2_2, -1_3.9_5_5_4], [-1_2.8_7_3_2, -1_3.9_3_5_2, -1_4.3_5_6_3], [-1_2.9_4_3_8, -1_3.8_2_2_6, -1_4.2_5_1_3]],
             [[-1_2.5_1_3_4, -1_3.4_6_8_6, -1_4.4_9_1_5], [-1_2.8_6_6_9, -1_4.4_3_4_3, -1_4.7_7_5_8], [-1_3.2_5_2_3, -1_4.5_8_1_9, -1_5.0_6_9_4]],
         ]				).to(__lowerCamelCase				)
     self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]							,       __lowerCamelCase							,       atol=1E-4				)				)



   @slow
   def    __UpperCAmelCase    (       self    :	List[Any]				):



     """simple docstring"""
     # only resize + normalize
     _snake_case   	=      SegformerImageProcessor(
         image_scale=(5_1_2, 5_1_2)							,       keep_ratio=__lowerCamelCase							,       align=__lowerCamelCase							,       do_random_crop=__lowerCamelCase				)
     _snake_case   	=      SegformerForSemanticSegmentation.from_pretrained(
         '''nvidia/segformer-b1-finetuned-cityscapes-1024-1024'''				).to(__lowerCamelCase				)

     _snake_case   	=      prepare_img()
     _snake_case   	=      image_processor(images=__lowerCamelCase							,       return_tensors='''pt'''				)
     _snake_case   	=      encoded_inputs.pixel_values.to(__lowerCamelCase				)

     with torch.no_grad():
       _snake_case   	=      model(__lowerCamelCase				)

     _snake_case   	=      torch.Size((1, model.config.num_labels, 1_2_8, 1_2_8)				)
     self.assertEqual(outputs.logits.shape							,       __lowerCamelCase				)

     _snake_case   	=      torch.tensor(
         [
             [[-1_3.5_7_4_8, -1_3.9_1_1_1, -1_2.6_5_0_0], [-1_4.3_5_0_0, -1_5.3_6_8_3, -1_4.2_3_2_8], [-1_4.7_5_3_2, -1_6.0_4_2_4, -1_5.6_0_8_7]],
             [[-1_7.1_6_5_1, -1_5.8_7_2_5, -1_2.9_6_5_3], [-1_7.2_5_8_0, -1_7.3_7_1_8, -1_4.8_2_2_3], [-1_6.6_0_5_8, -1_6.8_7_8_3, -1_6.7_4_5_2]],
             [[-3.6_4_5_6, -3.0_2_0_9, -1.4_2_0_3], [-3.0_7_9_7, -3.1_9_5_9, -2.0_0_0_0], [-1.8_7_5_7, -1.9_2_1_7, -1.6_9_9_7]],
         ]				).to(__lowerCamelCase				)
     self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]							,       __lowerCamelCase							,       atol=1E-1				)				)





   @slow
   def    __UpperCAmelCase    (       self    :	Union[str, Any]				):



     """simple docstring"""
     # only resize + normalize
     _snake_case   	=      SegformerImageProcessor(
         image_scale=(5_1_2, 5_1_2)							,       keep_ratio=__lowerCamelCase							,       align=__lowerCamelCase							,       do_random_crop=__lowerCamelCase				)
     _snake_case   	=      SegformerForSemanticSegmentation.from_pretrained('''nvidia/segformer-b0-finetuned-ade-512-512'''				).to(
         __lowerCamelCase				)

     _snake_case   	=      prepare_img()
     _snake_case   	=      image_processor(images=__lowerCamelCase							,       return_tensors='''pt'''				)
     _snake_case   	=      encoded_inputs.pixel_values.to(__lowerCamelCase				)

     with torch.no_grad():
       _snake_case   	=      model(__lowerCamelCase				)

     _snake_case   	=      outputs.logits.detach().cpu()

     _snake_case   	=      image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase							,       target_sizes=[(5_0_0, 3_0_0)]				)
     _snake_case   	=      torch.Size((5_0_0, 3_0_0)				)
     self.assertEqual(segmentation[0].shape							,       __lowerCamelCase				)

     _snake_case   	=      image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase				)
     _snake_case   	=      torch.Size((1_2_8, 1_2_8)				)
     self.assertEqual(segmentation[0].shape							,       __lowerCamelCase				)











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


lowercase__							:	Optional[Any]											=		logging.get_logger(__name__)

# General docstring
lowercase__							:	Optional[int]											=		"PoolFormerConfig"

# Base docstring
lowercase__							:	Optional[Any]											=		"sail/poolformer_s12"
lowercase__							:	Union[str, Any]											=		[1, 512, 7, 7]

# Image classification docstring
lowercase__							:	List[str]											=		"sail/poolformer_s12"
lowercase__							:	Dict											=		"tabby, tabby cat"

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


def 		lowerCamelCase__       (						_A	,      _A = 0.0	,      _A = False				):







	'''simple docstring'''





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







						'''simple docstring'''





						def __init__(		self				:	Optional[Any]			,      __lowercase				:	Optional[float] = None		):






							"""simple docstring"""






							super().__init__()
							snake_case_						       =	drop_prob






						def 			snake_case__		(		self				:	List[str]			,      __lowercase				:	torch.Tensor		):






							"""simple docstring"""






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






						def 			snake_case__		(		self				:	Any		):






							"""simple docstring"""






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







class 						UpperCAmelCase  (							nn.Module     ):







						'''simple docstring'''





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






							"""simple docstring"""






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

							snake_case_						       =	nn.Convad(__lowercase			,      __lowercase			,      kernel_size=__lowercase			,      stride=__lowercase			,      padding=__lowercase		)
							snake_case_						       =	norm_layer(__lowercase		) if norm_layer else nn.Identity()






						def 			snake_case__		(		self				:	Any			,      __lowercase				:	Optional[Any]		):






							"""simple docstring"""






							snake_case_						       =	self.projection(__lowercase		)
							snake_case_						       =	self.norm(__lowercase		)
							return embeddings







class 						UpperCAmelCase  (							nn.GroupNorm     ):







						'''simple docstring'''





						def __init__(		self				:	int			,      __lowercase				:	List[str]			,      **__lowercase				:	Union[str, Any]		):






							"""simple docstring"""






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







class 						UpperCAmelCase  (							nn.Module     ):







						'''simple docstring'''





						def __init__(		self				:	Dict			,      __lowercase				:	List[str]		):






							"""simple docstring"""






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






						def 			snake_case__		(		self				:	int			,      __lowercase				:	Tuple		):






							"""simple docstring"""






							return self.pool(__lowercase		) - hidden_states







class 						UpperCAmelCase  (							nn.Module     ):







						'''simple docstring'''





						def __init__(		self				:	str			,      __lowercase				:	List[Any]			,      __lowercase				:	str			,      __lowercase				:	Optional[Any]			,      __lowercase				:	Optional[Any]		):






							"""simple docstring"""






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






						def 			snake_case__		(		self				:	Union[str, Any]			,      __lowercase				:	int		):






							"""simple docstring"""






							snake_case_						       =	self.conva(__lowercase		)
							snake_case_						       =	self.act_fn(__lowercase		)
							snake_case_						       =	self.drop(__lowercase		)
							snake_case_						       =	self.conva(__lowercase		)
							snake_case_						       =	self.drop(__lowercase		)

							return hidden_states







class 						UpperCAmelCase  (							nn.Module     ):







						'''simple docstring'''





						def __init__(		self				:	int			,      __lowercase				:	List[str]			,      __lowercase				:	Optional[Any]			,      __lowercase				:	str			,      __lowercase				:	Optional[Any]			,      __lowercase				:	List[str]			,      __lowercase				:	int		):






							"""simple docstring"""






							super().__init__()
							snake_case_						       =	PoolFormerPooling(__lowercase		)
							snake_case_						       =	PoolFormerOutput(__lowercase			,      __lowercase			,      __lowercase			,      __lowercase		)
							snake_case_						       =	PoolFormerGroupNorm(__lowercase		)
							snake_case_						       =	PoolFormerGroupNorm(__lowercase		)

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






						def 			snake_case__		(		self				:	Union[str, Any]			,      __lowercase				:	Union[str, Any]		):






							"""simple docstring"""






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

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

								snake_case_						       =	(output,) + outputs
								return outputs

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

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

								snake_case_						       =	(output,) + outputs
								return outputs







class 						UpperCAmelCase  (							nn.Module     ):







						'''simple docstring'''





						def __init__(		self				:	List[Any]			,      __lowercase				:	Optional[Any]		):






							"""simple docstring"""






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

							# patch embeddings
							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]			,      )		)
							snake_case_						       =	nn.ModuleList(__lowercase		)

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

							snake_case_						       =	nn.ModuleList(__lowercase		)






						def 			snake_case__		(		self				:	List[str]			,      __lowercase				:	List[Any]			,      __lowercase				:	int=False			,      __lowercase				:	Tuple=True		):






							"""simple docstring"""






							snake_case_						       =	() if output_hidden_states else None

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

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







class 						UpperCAmelCase  (							UpperCAmelCase__     ):







						'''simple docstring'''





						lowerCAmelCase_							=  PoolFormerConfig
						lowerCAmelCase_							=  '''poolformer'''
						lowerCAmelCase_							=  '''pixel_values'''
						lowerCAmelCase_							=  True






						def 			snake_case__		(		self				:	Dict			,      __lowercase				:	Optional[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 			snake_case__		(		self				:	str			,      __lowercase				:	Any			,      __lowercase				:	Tuple=False		):






							"""simple docstring"""






							if isinstance(__lowercase			,      __lowercase		):
								snake_case_						       =	value







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

lowercase__							:	List[str]											=		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__     ):







						'''simple docstring'''





						def __init__(		self				:	Any			,      __lowercase				:	Any		):






							"""simple docstring"""






							super().__init__(__lowercase		)
							snake_case_						       =	config

							snake_case_						       =	PoolFormerEncoder(__lowercase		)

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






						def 			snake_case__		(		self				:	List[Any]		):






							"""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 			snake_case__		(		self				:	Optional[Any]			,      __lowercase				:	Optional[torch.FloatTensor] = None			,      __lowercase				:	Optional[bool] = None			,      __lowercase				:	Optional[bool] = None			,      ):






							"""simple docstring"""






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

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







						'''simple docstring'''





						def __init__(		self				:	Dict			,      __lowercase				:	Tuple		):






							"""simple docstring"""






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






						def 			snake_case__		(		self				:	Dict			,      __lowercase				:	List[str]		):






							"""simple docstring"""






							snake_case_						       =	self.dense(__lowercase		)
							return output







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







						'''simple docstring'''





						def __init__(		self				:	Optional[int]			,      __lowercase				:	int		):






							"""simple docstring"""






							super().__init__(__lowercase		)
							snake_case_						       =	config.num_labels
							snake_case_						       =	PoolFormerModel(__lowercase		)

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






							"""simple docstring"""






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

							snake_case_						       =	self.poolformer(
							    __lowercase			,      output_hidden_states=__lowercase			,      return_dict=__lowercase			,      )

							snake_case_						       =	outputs[0]

							snake_case_						       =	self.classifier(self.norm(__lowercase		).mean([-2, -1]		)		)

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

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

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

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