


"""simple docstring"""





import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
 import torch
 from torch import nn

 from transformers import (
     MODEL_MAPPING,
     BeitForImageClassification,
     BeitForMaskedImageModeling,
     BeitForSemanticSegmentation,
     BeitModel,
 )
 from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 import PIL
 from PIL import Image

 from transformers import BeitImageProcessor



class 		snake_case_:






     def __init__(	self  :					Optional[int]					,				UpperCamelCase_  :					List[Any]					,				UpperCamelCase_  :					List[str]=1_0_0					,				UpperCamelCase_  :					List[str]=1_3					,				UpperCamelCase_  :					Optional[int]=3_0					,				UpperCamelCase_  :					List[str]=2					,				UpperCamelCase_  :					Union[str, Any]=3					,				UpperCamelCase_  :					str=True					,				UpperCamelCase_  :					int=True					,				UpperCamelCase_  :					Optional[Any]=3_2					,				UpperCamelCase_  :					Union[str, Any]=4					,				UpperCamelCase_  :					int=4					,				UpperCamelCase_  :					Any=3_7					,				UpperCamelCase_  :					str="gelu"					,				UpperCamelCase_  :					Optional[Any]=0.1					,				UpperCamelCase_  :					Optional[int]=0.1					,				UpperCamelCase_  :					int=1_0					,				UpperCamelCase_  :					Tuple=0.02					,				UpperCamelCase_  :					Tuple=3					,				UpperCamelCase_  :					Optional[int]=None					,				UpperCamelCase_  :					Dict=[0, 1, 2, 3]					,				):
           lowerCAmelCase					:	Dict       =   parent
           lowerCAmelCase					:	List[Any]       =   1_0_0
           lowerCAmelCase					:	List[str]       =   batch_size
           lowerCAmelCase					:	Union[str, Any]       =   image_size
           lowerCAmelCase					:	Dict       =   patch_size
           lowerCAmelCase					:	Any       =   num_channels
           lowerCAmelCase					:	List[Any]       =   is_training
           lowerCAmelCase					:	Union[str, Any]       =   use_labels
           lowerCAmelCase					:	Tuple       =   hidden_size
           lowerCAmelCase					:	Dict       =   num_hidden_layers
           lowerCAmelCase					:	Union[str, Any]       =   num_attention_heads
           lowerCAmelCase					:	Optional[Any]       =   intermediate_size
           lowerCAmelCase					:	Dict       =   hidden_act
           lowerCAmelCase					:	Tuple       =   hidden_dropout_prob
           lowerCAmelCase					:	List[Any]       =   attention_probs_dropout_prob
           lowerCAmelCase					:	Union[str, Any]       =   type_sequence_label_size
           lowerCAmelCase					:	Any       =   initializer_range
           lowerCAmelCase					:	Optional[Any]       =   scope
           lowerCAmelCase					:	Tuple       =   out_indices
           lowerCAmelCase					:	List[str]       =   num_labels

           # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
           lowerCAmelCase					:	List[Any]       =   (image_size // patch_size) ** 2
           lowerCAmelCase					:	List[Any]       =   num_patches + 1






     def      lowerCamelCase__     (	self  :					str				):
           lowerCAmelCase					:	Dict       =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

           lowerCAmelCase					:	List[Any]       =   None
           lowerCAmelCase					:	Optional[int]       =   None
           if self.use_labels:
                 lowerCAmelCase					:	str       =   ids_tensor([self.batch_size]					,				self.type_sequence_label_size				)
                 lowerCAmelCase					:	List[Any]       =   ids_tensor([self.batch_size, self.image_size, self.image_size]					,				self.num_labels				)

           lowerCAmelCase					:	Tuple       =   self.get_config()

           return config, pixel_values, labels, pixel_labels






     def      lowerCamelCase__     (	self  :					Any				):
           return BeitConfig(
               vocab_size=self.vocab_size					,				image_size=self.image_size					,				patch_size=self.patch_size					,				num_channels=self.num_channels					,				hidden_size=self.hidden_size					,				num_hidden_layers=self.num_hidden_layers					,				num_attention_heads=self.num_attention_heads					,				intermediate_size=self.intermediate_size					,				hidden_act=self.hidden_act					,				hidden_dropout_prob=self.hidden_dropout_prob					,				attention_probs_dropout_prob=self.attention_probs_dropout_prob					,				is_decoder=UpperCamelCase_					,				initializer_range=self.initializer_range					,				out_indices=self.out_indices					,				)






     def      lowerCamelCase__     (	self  :					List[Any]					,				UpperCamelCase_  :					List[Any]					,				UpperCamelCase_  :					Tuple					,				UpperCamelCase_  :					Optional[Any]					,				UpperCamelCase_  :					str				):
           lowerCAmelCase					:	Union[str, Any]       =   BeitModel(config=UpperCamelCase_				)
           model.to(UpperCamelCase_				)
           model.eval()
           lowerCAmelCase					:	List[str]       =   model(UpperCamelCase_				)
           self.parent.assertEqual(result.last_hidden_state.shape					,				(self.batch_size, self.seq_length, self.hidden_size)				)






     def      lowerCamelCase__     (	self  :					int					,				UpperCamelCase_  :					str					,				UpperCamelCase_  :					str					,				UpperCamelCase_  :					Dict					,				UpperCamelCase_  :					Tuple				):
           lowerCAmelCase					:	List[Any]       =   BeitForMaskedImageModeling(config=UpperCamelCase_				)
           model.to(UpperCamelCase_				)
           model.eval()
           lowerCAmelCase					:	Union[str, Any]       =   model(UpperCamelCase_				)
           self.parent.assertEqual(result.logits.shape					,				(self.batch_size, self.seq_length - 1, self.vocab_size)				)






     def      lowerCamelCase__     (	self  :					Tuple					,				UpperCamelCase_  :					Dict					,				UpperCamelCase_  :					Optional[Any]					,				UpperCamelCase_  :					List[str]					,				UpperCamelCase_  :					int				):
           lowerCAmelCase					:	Union[str, Any]       =   self.type_sequence_label_size
           lowerCAmelCase					:	Union[str, Any]       =   BeitForImageClassification(UpperCamelCase_				)
           model.to(UpperCamelCase_				)
           model.eval()
           lowerCAmelCase					:	Optional[Any]       =   model(UpperCamelCase_					,				labels=UpperCamelCase_				)
           self.parent.assertEqual(result.logits.shape					,				(self.batch_size, self.type_sequence_label_size)				)

           # test greyscale images
           lowerCAmelCase					:	List[str]       =   1
           lowerCAmelCase					:	str       =   BeitForImageClassification(UpperCamelCase_				)
           model.to(UpperCamelCase_				)
           model.eval()

           lowerCAmelCase					:	Any       =   floats_tensor([self.batch_size, 1, self.image_size, self.image_size]				)
           lowerCAmelCase					:	Tuple       =   model(UpperCamelCase_					,				labels=UpperCamelCase_				)
           self.parent.assertEqual(result.logits.shape					,				(self.batch_size, self.type_sequence_label_size)				)






     def      lowerCamelCase__     (	self  :					Optional[int]					,				UpperCamelCase_  :					Union[str, Any]					,				UpperCamelCase_  :					Tuple					,				UpperCamelCase_  :					Optional[Any]					,				UpperCamelCase_  :					int				):
           lowerCAmelCase					:	List[str]       =   self.num_labels
           lowerCAmelCase					:	int       =   BeitForSemanticSegmentation(UpperCamelCase_				)
           model.to(UpperCamelCase_				)
           model.eval()
           lowerCAmelCase					:	int       =   model(UpperCamelCase_				)
           self.parent.assertEqual(
               result.logits.shape					,				(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)				)
           lowerCAmelCase					:	Tuple       =   model(UpperCamelCase_					,				labels=UpperCamelCase_				)
           self.parent.assertEqual(
               result.logits.shape					,				(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)				)






     def      lowerCamelCase__     (	self  :					Optional[int]				):
           lowerCAmelCase					:	Optional[Any]       =   self.prepare_config_and_inputs()
           lowerCAmelCase,      lowerCAmelCase,      lowerCAmelCase,      lowerCAmelCase					:	Any       =   config_and_inputs
           lowerCAmelCase					:	Optional[Any]       =   {'''pixel_values''': pixel_values}
           return config, inputs_dict



@require_torch
class 		snake_case_(       a__							,					a__							,					unittest.TestCase			):
     __UpperCamelCase  =					(
         (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
         if is_torch_available()
         else ()
     )
     __UpperCamelCase  =					(
         {
             '''feature-extraction''': BeitModel,
             '''image-classification''': BeitForImageClassification,
             '''image-segmentation''': BeitForSemanticSegmentation,
         }
         if is_torch_available()
         else {}
     )

     __UpperCamelCase  =					False
     __UpperCamelCase  =					False
     __UpperCamelCase  =					False






     def      lowerCamelCase__     (	self  :					Any				):
           lowerCAmelCase					:	str       =   BeitModelTester(self				)
           lowerCAmelCase					:	Tuple       =   ConfigTester(self					,				config_class=UpperCamelCase_					,				has_text_modality=UpperCamelCase_					,				hidden_size=3_7				)






     def      lowerCamelCase__     (	self  :					Optional[Any]				):
           self.config_tester.run_common_tests()






     @unittest.skip(reason='''BEiT does not use inputs_embeds'''				)
     def      lowerCamelCase__     (	self  :					str				):
           pass






     @require_torch_multi_gpu
     @unittest.skip(reason='''BEiT has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''				)
     def      lowerCamelCase__     (	self  :					Optional[int]				):
           pass






     def      lowerCamelCase__     (	self  :					List[Any]				):
           lowerCAmelCase,      lowerCAmelCase					:	List[Any]       =   self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 lowerCAmelCase					:	Union[str, Any]       =   model_class(UpperCamelCase_				)
                 self.assertIsInstance(model.get_input_embeddings()					,				(nn.Module)				)
                 lowerCAmelCase					:	Dict       =   model.get_output_embeddings()
                 self.assertTrue(x is None or isinstance(UpperCamelCase_					,				nn.Linear				)				)






     def      lowerCamelCase__     (	self  :					Dict				):
           lowerCAmelCase,      lowerCAmelCase					:	Tuple       =   self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 lowerCAmelCase					:	Any       =   model_class(UpperCamelCase_				)
                 lowerCAmelCase					:	Union[str, Any]       =   inspect.signature(model.forward				)
                 # signature.parameters is an OrderedDict => so arg_names order is deterministic
                 lowerCAmelCase					:	List[Any]       =   [*signature.parameters.keys()]

                 lowerCAmelCase					:	Dict       =   ['''pixel_values''']
                 self.assertListEqual(arg_names[:1]					,				UpperCamelCase_				)






     def      lowerCamelCase__     (	self  :					Any				):
           lowerCAmelCase					:	List[str]       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase_				)






     def      lowerCamelCase__     (	self  :					Dict				):
           lowerCAmelCase					:	str       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase_				)






     def      lowerCamelCase__     (	self  :					Optional[Any]				):
           lowerCAmelCase					:	Optional[Any]       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCamelCase_				)






     def      lowerCamelCase__     (	self  :					int				):
           lowerCAmelCase					:	Union[str, Any]       =   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_semantic_segmentation(*UpperCamelCase_				)






     def      lowerCamelCase__     (	self  :					Tuple				):
           if not self.model_tester.is_training:
                 return

           lowerCAmelCase,      lowerCAmelCase					:	Dict       =   self.model_tester.prepare_config_and_inputs_for_common()
           lowerCAmelCase					:	Optional[int]       =   True

           for model_class in self.all_model_classes:
                 # we don't test BeitForMaskedImageModeling
                 if model_class in [*get_values(UpperCamelCase_				), BeitForMaskedImageModeling]:
                       continue

                 lowerCAmelCase					:	Tuple       =   model_class(UpperCamelCase_				)
                 model.to(UpperCamelCase_				)
                 model.train()
                 lowerCAmelCase					:	str       =   self._prepare_for_class(UpperCamelCase_					,				UpperCamelCase_					,				return_labels=UpperCamelCase_				)
                 lowerCAmelCase					:	List[Any]       =   model(**UpperCamelCase_				).loss
                 loss.backward()






     def      lowerCamelCase__     (	self  :					Optional[Any]				):
           lowerCAmelCase,      lowerCAmelCase					:	List[Any]       =   self.model_tester.prepare_config_and_inputs_for_common()
           if not self.model_tester.is_training:
                 return

           lowerCAmelCase					:	Any       =   False
           lowerCAmelCase					:	Dict       =   True

           for model_class in self.all_model_classes:
                 # we don't test BeitForMaskedImageModeling
                 if (
                     model_class in [*get_values(UpperCamelCase_				), BeitForMaskedImageModeling]
                     or not model_class.supports_gradient_checkpointing
                 ):
                       continue

                 lowerCAmelCase					:	str       =   model_class(UpperCamelCase_				)
                 model.gradient_checkpointing_enable()
                 model.to(UpperCamelCase_				)
                 model.train()
                 lowerCAmelCase					:	Tuple       =   self._prepare_for_class(UpperCamelCase_					,				UpperCamelCase_					,				return_labels=UpperCamelCase_				)
                 lowerCAmelCase					:	Union[str, Any]       =   model(**UpperCamelCase_				).loss
                 loss.backward()






     def      lowerCamelCase__     (	self  :					Union[str, Any]				):
           lowerCAmelCase,      lowerCAmelCase					:	int       =   self.model_tester.prepare_config_and_inputs_for_common()

           lowerCAmelCase					:	Any       =   _config_zero_init(UpperCamelCase_				)
           for model_class in self.all_model_classes:
                 lowerCAmelCase					:	str       =   model_class(config=UpperCamelCase_				)
                 for name, param in model.named_parameters():
                       # we skip lambda parameters as these require special initial values
                       # determined by config.layer_scale_init_value
                       if "lambda" in name:
                             continue
                       if param.requires_grad:
                             self.assertIn(
                                 ((param.data.mean() * 1E9).round() / 1E9).item()					,				[0.0, 1.0]					,				msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''					,				)






     @slow
     def      lowerCamelCase__     (	self  :					Any				):
           for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCAmelCase					:	Tuple       =   BeitModel.from_pretrained(UpperCamelCase_				)
                 self.assertIsNotNone(UpperCamelCase_				)



def      _snake_case						(						):
      lowerCAmelCase					:	str       =   Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''			)
      return image






@require_torch
@require_vision
class 		snake_case_(       unittest.TestCase			):






     @cached_property
     def      lowerCamelCase__     (	self  :					Dict				):
           return BeitImageProcessor.from_pretrained('''microsoft/beit-base-patch16-224'''				) if is_vision_available() else None






     @slow
     def      lowerCamelCase__     (	self  :					List[str]				):
           lowerCAmelCase					:	Union[str, Any]       =   BeitForMaskedImageModeling.from_pretrained('''microsoft/beit-base-patch16-224-pt22k'''				).to(UpperCamelCase_				)

           lowerCAmelCase					:	Optional[Any]       =   self.default_image_processor
           lowerCAmelCase					:	Union[str, Any]       =   prepare_img()
           lowerCAmelCase					:	int       =   image_processor(images=UpperCamelCase_					,				return_tensors='''pt'''				).pixel_values.to(UpperCamelCase_				)

           # prepare bool_masked_pos
           lowerCAmelCase					:	List[Any]       =   torch.ones((1, 1_9_6)					,				dtype=torch.bool				).to(UpperCamelCase_				)

           # forward pass
           with torch.no_grad():
                 lowerCAmelCase					:	List[Any]       =   model(pixel_values=UpperCamelCase_					,				bool_masked_pos=UpperCamelCase_				)
           lowerCAmelCase					:	List[str]       =   outputs.logits

           # verify the logits
           lowerCAmelCase					:	List[Any]       =   torch.Size((1, 1_9_6, 8_1_9_2)				)
           self.assertEqual(logits.shape					,				UpperCamelCase_				)

           lowerCAmelCase					:	Optional[int]       =   torch.tensor(
               [[-3.2_437, 0.5_072, -13.9_174], [-3.2_456, 0.4_948, -13.9_401], [-3.2_033, 0.5_121, -13.8_550]]				).to(UpperCamelCase_				)

           self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]					,				UpperCamelCase_					,				atol=1E-2				)				)






     @slow
     def      lowerCamelCase__     (	self  :					int				):
           lowerCAmelCase					:	Dict       =   BeitForImageClassification.from_pretrained('''microsoft/beit-base-patch16-224'''				).to(UpperCamelCase_				)

           lowerCAmelCase					:	Union[str, Any]       =   self.default_image_processor
           lowerCAmelCase					:	Union[str, Any]       =   prepare_img()
           lowerCAmelCase					:	int       =   image_processor(images=UpperCamelCase_					,				return_tensors='''pt'''				).to(UpperCamelCase_				)

           # forward pass
           with torch.no_grad():
                 lowerCAmelCase					:	int       =   model(**UpperCamelCase_				)
           lowerCAmelCase					:	List[Any]       =   outputs.logits

           # verify the logits
           lowerCAmelCase					:	List[str]       =   torch.Size((1, 1_0_0_0)				)
           self.assertEqual(logits.shape					,				UpperCamelCase_				)

           lowerCAmelCase					:	Optional[int]       =   torch.tensor([-1.2_385, -1.0_987, -1.0_108]				).to(UpperCamelCase_				)

           self.assertTrue(torch.allclose(logits[0, :3]					,				UpperCamelCase_					,				atol=1E-4				)				)

           lowerCAmelCase					:	Optional[int]       =   2_8_1
           self.assertEqual(logits.argmax(-1				).item()					,				UpperCamelCase_				)






     @slow
     def      lowerCamelCase__     (	self  :					Union[str, Any]				):
           lowerCAmelCase					:	Tuple       =   BeitForImageClassification.from_pretrained('''microsoft/beit-large-patch16-224-pt22k-ft22k'''				).to(
               UpperCamelCase_				)

           lowerCAmelCase					:	Optional[Any]       =   self.default_image_processor
           lowerCAmelCase					:	Optional[int]       =   prepare_img()
           lowerCAmelCase					:	int       =   image_processor(images=UpperCamelCase_					,				return_tensors='''pt'''				).to(UpperCamelCase_				)

           # forward pass
           with torch.no_grad():
                 lowerCAmelCase					:	Optional[int]       =   model(**UpperCamelCase_				)
           lowerCAmelCase					:	Tuple       =   outputs.logits

           # verify the logits
           lowerCAmelCase					:	Tuple       =   torch.Size((1, 2_1_8_4_1)				)
           self.assertEqual(logits.shape					,				UpperCamelCase_				)

           lowerCAmelCase					:	Optional[Any]       =   torch.tensor([1.6_881, -0.2_787, 0.5_901]				).to(UpperCamelCase_				)

           self.assertTrue(torch.allclose(logits[0, :3]					,				UpperCamelCase_					,				atol=1E-4				)				)

           lowerCAmelCase					:	Optional[int]       =   2_3_9_6
           self.assertEqual(logits.argmax(-1				).item()					,				UpperCamelCase_				)






     @slow
     def      lowerCamelCase__     (	self  :					Any				):
           lowerCAmelCase					:	Optional[Any]       =   BeitForSemanticSegmentation.from_pretrained('''microsoft/beit-base-finetuned-ade-640-640'''				)
           lowerCAmelCase					:	str       =   model.to(UpperCamelCase_				)

           lowerCAmelCase					:	Union[str, Any]       =   BeitImageProcessor(do_resize=UpperCamelCase_					,				size=6_4_0					,				do_center_crop=UpperCamelCase_				)

           lowerCAmelCase					:	Any       =   load_dataset('''hf-internal-testing/fixtures_ade20k'''					,				split='''test'''				)
           lowerCAmelCase					:	int       =   Image.open(ds[0]['''file''']				)
           lowerCAmelCase					:	Tuple       =   image_processor(images=UpperCamelCase_					,				return_tensors='''pt'''				).to(UpperCamelCase_				)

           # forward pass
           with torch.no_grad():
                 lowerCAmelCase					:	List[Any]       =   model(**UpperCamelCase_				)
           lowerCAmelCase					:	Any       =   outputs.logits

           # verify the logits
           lowerCAmelCase					:	List[str]       =   torch.Size((1, 1_5_0, 1_6_0, 1_6_0)				)
           self.assertEqual(logits.shape					,				UpperCamelCase_				)

           lowerCAmelCase					:	List[str]       =   version.parse(PIL.__version__				) < version.parse('''9.0.0'''				)

           if is_pillow_less_than_a:
                 lowerCAmelCase					:	Any       =   torch.tensor(
                     [
                         [[-4.9_225, -2.3_954, -3.0_522], [-2.8_822, -1.0_046, -1.7_561], [-2.9_549, -1.3_228, -2.1_347]],
                         [[-5.8_168, -3.4_129, -4.0_778], [-3.8_651, -2.2_214, -3.0_277], [-3.8_356, -2.4_643, -3.3_535]],
                         [[-0.0_078, 3.9_952, 4.0_754], [2.9_856, 4.6_944, 5.0_035], [3.2_413, 4.7_813, 4.9_969]],
                     ]					,				device=UpperCamelCase_					,				)
           else:
                 lowerCAmelCase					:	Tuple       =   torch.tensor(
                     [
                         [[-4.8_960, -2.3_688, -3.0_355], [-2.8_478, -0.9_836, -1.7_418], [-2.9_449, -1.3_332, -2.1_456]],
                         [[-5.8_081, -3.4_124, -4.1_006], [-3.8_561, -2.2_081, -3.0_323], [-3.8_365, -2.4_601, -3.3_669]],
                         [[-0.0_309, 3.9_868, 4.0_540], [2.9_640, 4.6_877, 4.9_976], [3.2_081, 4.7_690, 4.9_942]],
                     ]					,				device=UpperCamelCase_					,				)

           self.assertTrue(torch.allclose(logits[0, :3, :3, :3]					,				UpperCamelCase_					,				atol=1E-4				)				)






     @slow
     def      lowerCamelCase__     (	self  :					Dict				):
           lowerCAmelCase					:	Union[str, Any]       =   BeitForSemanticSegmentation.from_pretrained('''microsoft/beit-base-finetuned-ade-640-640'''				)
           lowerCAmelCase					:	Optional[int]       =   model.to(UpperCamelCase_				)

           lowerCAmelCase					:	Dict       =   BeitImageProcessor(do_resize=UpperCamelCase_					,				size=6_4_0					,				do_center_crop=UpperCamelCase_				)

           lowerCAmelCase					:	Dict       =   load_dataset('''hf-internal-testing/fixtures_ade20k'''					,				split='''test'''				)
           lowerCAmelCase					:	str       =   Image.open(ds[0]['''file''']				)
           lowerCAmelCase					:	List[str]       =   image_processor(images=UpperCamelCase_					,				return_tensors='''pt'''				).to(UpperCamelCase_				)

           # forward pass
           with torch.no_grad():
                 lowerCAmelCase					:	Optional[Any]       =   model(**UpperCamelCase_				)

           lowerCAmelCase					:	List[str]       =   outputs.logits.detach().cpu()

           lowerCAmelCase					:	str       =   image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase_					,				target_sizes=[(5_0_0, 3_0_0)]				)
           lowerCAmelCase					:	List[Any]       =   torch.Size((5_0_0, 3_0_0)				)
           self.assertEqual(segmentation[0].shape					,				UpperCamelCase_				)

           lowerCAmelCase					:	List[str]       =   image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase_				)
           lowerCAmelCase					:	str       =   torch.Size((1_6_0, 1_6_0)				)
           self.assertEqual(segmentation[0].shape					,				UpperCamelCase_				)






"""simple docstring"""





from torch import nn



class 		snake_case_(       nn.Module			):






     def __init__(	self  :					int					,				UpperCamelCase_  :					int					,				UpperCamelCase_  :					int				):
           super().__init__()
           lowerCAmelCase					:	str       =   class_size
           lowerCAmelCase					:	Dict       =   embed_size
           # self.mlp1 = nn.Linear(embed_size, embed_size)
           # self.mlp2 = (nn.Linear(embed_size, class_size))
           lowerCAmelCase					:	Any       =   nn.Linear(UpperCamelCase_					,				UpperCamelCase_				)






     def      lowerCamelCase__     (	self  :					List[Any]					,				UpperCamelCase_  :					Tuple				):
           # hidden_state = nn.functional.relu(self.mlp1(hidden_state))
           # hidden_state = self.mlp2(hidden_state)
           lowerCAmelCase					:	int       =   self.mlp(UpperCamelCase_				)
           return logits


