





'''simple docstring'''




from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)
@flax.struct.dataclass
class 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






     snake_case_						      =   42




@flax_register_to_config
class 			_lowerCAmelCase		(			nn.Module							,     A__							,     A__	):






     """simple docstring"""






     snake_case_						      =   32
     snake_case_						      =   4
     snake_case_						      =   4
     snake_case_						      =   (
         "CrossAttnDownBlock2D",
         "CrossAttnDownBlock2D",
         "CrossAttnDownBlock2D",
         "DownBlock2D",
     )
     snake_case_						      =   ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
     snake_case_						      =   False
     snake_case_						      =   (320, 640, 1280, 1280)
     snake_case_						      =   2
     snake_case_						      =   8
     snake_case_						      =   None
     snake_case_						      =   1280
     snake_case_						      =   0.0
     snake_case_						      =   False
     snake_case_						      =   jnp.floataa
     snake_case_						      =   True
     snake_case_						      =   0
     snake_case_						      =   False



     def      lowerCAmelCase			(   self							:			str   ,       __snake_case							:			jax.random.KeyArray						)->	FrozenDict:
        snake_case 						=   (1, self.in_channels, self.sample_size, self.sample_size)
        snake_case 						=   jnp.zeros(lowerCAmelCase__   ,       dtype=jnp.floataa						)
        snake_case 						=   jnp.ones((1,)   ,       dtype=jnp.intaa						)
        snake_case 						=   jnp.zeros((1, 1, self.cross_attention_dim)   ,       dtype=jnp.floataa						)

        snake_case						,    snake_case 						=   jax.random.split(lowerCAmelCase__						)
        snake_case 						=   {"""params""": params_rng, """dropout""": dropout_rng}

        return self.init(lowerCAmelCase__   ,       lowerCAmelCase__   ,       lowerCAmelCase__   ,       lowerCAmelCase__						)["params"]



     def      lowerCAmelCase			(   self							:			str						)->	Tuple:
        snake_case 						=   self.block_out_channels
        snake_case 						=   block_out_channels[0] * 4

        if self.num_attention_heads is not None:
           raise ValueError(
               """At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."""						)

        # If `num_attention_heads` is not defined (which is the case for most models)
        # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
        # The reason for this behavior is to correct for incorrectly named variables that were introduced
        # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
        # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
        # which is why we correct for the naming here.
        snake_case 						=   self.num_attention_heads or self.attention_head_dim

        # input
        snake_case 						=   nn.Conv(
            block_out_channels[0]   ,       kernel_size=(3, 3)   ,       strides=(1, 1)   ,       padding=((1, 1), (1, 1))   ,       dtype=self.dtype   ,       )

        # time
        snake_case 						=   FlaxTimesteps(
            block_out_channels[0]   ,       flip_sin_to_cos=self.flip_sin_to_cos   ,       freq_shift=self.config.freq_shift						)
        snake_case 						=   FlaxTimestepEmbedding(lowerCAmelCase__   ,       dtype=self.dtype						)

        snake_case 						=   self.only_cross_attention
        if isinstance(lowerCAmelCase__   ,       lowerCAmelCase__						):
           snake_case 						=   (only_cross_attention,) * len(self.down_block_types						)

        if isinstance(lowerCAmelCase__   ,       lowerCAmelCase__						):
           snake_case 						=   (num_attention_heads,) * len(self.down_block_types						)

        # down
        snake_case 						=   []
        snake_case 						=   block_out_channels[0]
        for i, down_block_type in enumerate(self.down_block_types						):
           snake_case 						=   output_channel
           snake_case 						=   block_out_channels[i]
           snake_case 						=   i == len(lowerCAmelCase__						) - 1

           if down_block_type == "CrossAttnDownBlock2D":
              snake_case 						=   FlaxCrossAttnDownBlockaD(
                  in_channels=lowerCAmelCase__   ,       out_channels=lowerCAmelCase__   ,       dropout=self.dropout   ,       num_layers=self.layers_per_block   ,       num_attention_heads=num_attention_heads[i]   ,       add_downsample=not is_final_block   ,       use_linear_projection=self.use_linear_projection   ,       only_cross_attention=only_cross_attention[i]   ,       use_memory_efficient_attention=self.use_memory_efficient_attention   ,       dtype=self.dtype   ,       )
           else:
              snake_case 						=   FlaxDownBlockaD(
                  in_channels=lowerCAmelCase__   ,       out_channels=lowerCAmelCase__   ,       dropout=self.dropout   ,       num_layers=self.layers_per_block   ,       add_downsample=not is_final_block   ,       dtype=self.dtype   ,       )

           down_blocks.append(lowerCAmelCase__						)
        snake_case 						=   down_blocks

        # mid
        snake_case 						=   FlaxUNetMidBlockaDCrossAttn(
            in_channels=block_out_channels[-1]   ,       dropout=self.dropout   ,       num_attention_heads=num_attention_heads[-1]   ,       use_linear_projection=self.use_linear_projection   ,       use_memory_efficient_attention=self.use_memory_efficient_attention   ,       dtype=self.dtype   ,       )

        # up
        snake_case 						=   []
        snake_case 						=   list(reversed(lowerCAmelCase__						)						)
        snake_case 						=   list(reversed(lowerCAmelCase__						)						)
        snake_case 						=   list(reversed(lowerCAmelCase__						)						)
        snake_case 						=   reversed_block_out_channels[0]
        for i, up_block_type in enumerate(self.up_block_types						):
           snake_case 						=   output_channel
           snake_case 						=   reversed_block_out_channels[i]
           snake_case 						=   reversed_block_out_channels[min(i + 1   ,       len(lowerCAmelCase__						) - 1						)]

           snake_case 						=   i == len(lowerCAmelCase__						) - 1

           if up_block_type == "CrossAttnUpBlock2D":
              snake_case 						=   FlaxCrossAttnUpBlockaD(
                  in_channels=lowerCAmelCase__   ,       out_channels=lowerCAmelCase__   ,       prev_output_channel=lowerCAmelCase__   ,       num_layers=self.layers_per_block + 1   ,       num_attention_heads=reversed_num_attention_heads[i]   ,       add_upsample=not is_final_block   ,       dropout=self.dropout   ,       use_linear_projection=self.use_linear_projection   ,       only_cross_attention=only_cross_attention[i]   ,       use_memory_efficient_attention=self.use_memory_efficient_attention   ,       dtype=self.dtype   ,       )
           else:
              snake_case 						=   FlaxUpBlockaD(
                  in_channels=lowerCAmelCase__   ,       out_channels=lowerCAmelCase__   ,       prev_output_channel=lowerCAmelCase__   ,       num_layers=self.layers_per_block + 1   ,       add_upsample=not is_final_block   ,       dropout=self.dropout   ,       dtype=self.dtype   ,       )

           up_blocks.append(lowerCAmelCase__						)
           snake_case 						=   output_channel
        snake_case 						=   up_blocks

        # out
        snake_case 						=   nn.GroupNorm(num_groups=32   ,       epsilon=1e-5						)
        snake_case 						=   nn.Conv(
            self.out_channels   ,       kernel_size=(3, 3)   ,       strides=(1, 1)   ,       padding=((1, 1), (1, 1))   ,       dtype=self.dtype   ,       )



     def __call__(   self							:			Tuple   ,       __snake_case							:			Dict   ,       __snake_case							:			List[Any]   ,       __snake_case							:			Dict   ,       __snake_case							:			Tuple=None   ,       __snake_case							:			Union[str, Any]=None   ,       __snake_case							:			bool = True   ,       __snake_case							:			bool = False   ,       )->	Union[FlaxUNetaDConditionOutput, Tuple]:




        if not isinstance(lowerCAmelCase__   ,       jnp.ndarray						):
           snake_case 						=   jnp.array([timesteps]   ,       dtype=jnp.intaa						)
        elif isinstance(lowerCAmelCase__   ,       jnp.ndarray						) and len(timesteps.shape						) == 0:
           snake_case 						=   timesteps.astype(dtype=jnp.floataa						)
           snake_case 						=   jnp.expand_dims(lowerCAmelCase__   ,       0						)

        snake_case 						=   self.time_proj(lowerCAmelCase__						)
        snake_case 						=   self.time_embedding(lowerCAmelCase__						)

        # 2. pre-process
        snake_case 						=   jnp.transpose(lowerCAmelCase__   ,       (0, 2, 3, 1)						)
        snake_case 						=   self.conv_in(lowerCAmelCase__						)

        # 3. down
        snake_case 						=   (sample,)
        for down_block in self.down_blocks:
           if isinstance(lowerCAmelCase__   ,       lowerCAmelCase__						):
              snake_case						,    snake_case 						=   down_block(lowerCAmelCase__   ,       lowerCAmelCase__   ,       lowerCAmelCase__   ,       deterministic=not train						)
           else:
              snake_case						,    snake_case 						=   down_block(lowerCAmelCase__   ,       lowerCAmelCase__   ,       deterministic=not train						)
           down_block_res_samples += res_samples

        if down_block_additional_residuals is not None:
           snake_case 						=   ()

           for down_block_res_sample, down_block_additional_residual in zip(
               lowerCAmelCase__   ,       lowerCAmelCase__						):
              down_block_res_sample += down_block_additional_residual
              new_down_block_res_samples += (down_block_res_sample,)

           snake_case 						=   new_down_block_res_samples

        # 4. mid
        snake_case 						=   self.mid_block(lowerCAmelCase__   ,       lowerCAmelCase__   ,       lowerCAmelCase__   ,       deterministic=not train						)

        if mid_block_additional_residual is not None:
           sample += mid_block_additional_residual

        # 5. up
        for up_block in self.up_blocks:
           snake_case 						=   down_block_res_samples[-(self.layers_per_block + 1) :]
           snake_case 						=   down_block_res_samples[: -(self.layers_per_block + 1)]
           if isinstance(lowerCAmelCase__   ,       lowerCAmelCase__						):
              snake_case 						=   up_block(
                  lowerCAmelCase__   ,       temb=lowerCAmelCase__   ,       encoder_hidden_states=lowerCAmelCase__   ,       res_hidden_states_tuple=lowerCAmelCase__   ,       deterministic=not train   ,       )
           else:
              snake_case 						=   up_block(lowerCAmelCase__   ,       temb=lowerCAmelCase__   ,       res_hidden_states_tuple=lowerCAmelCase__   ,       deterministic=not train						)

        # 6. post-process
        snake_case 						=   self.conv_norm_out(lowerCAmelCase__						)
        snake_case 						=   nn.silu(lowerCAmelCase__						)
        snake_case 						=   self.conv_out(lowerCAmelCase__						)
        snake_case 						=   jnp.transpose(lowerCAmelCase__   ,       (0, 3, 1, 2)						)

        if not return_dict:
           return (sample,)

        return FlaxUNetaDConditionOutput(sample=lowerCAmelCase__						)








'''simple docstring'''

import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

   if is_vision_available():
      from transformers import MaskaFormerImageProcessor

if is_vision_available():
   from PIL import Image
class 							__lowerCAmelCase	:






       """simple docstring"""





       def __init__(				self  :						Tuple			,					lowerCAmelCase__  :						List[Any]			,					lowerCAmelCase__  :						Tuple=2			,					lowerCAmelCase__  :						Optional[int]=True			,					lowerCAmelCase__  :						List[str]=False			,					lowerCAmelCase__  :						Tuple=10			,					lowerCAmelCase__  :						int=3			,					lowerCAmelCase__  :						Union[str, Any]=32 * 8			,					lowerCAmelCase__  :						Dict=32 * 8			,					lowerCAmelCase__  :						Dict=4			,					lowerCAmelCase__  :						Tuple=64			,					)       ->	List[Any]:

            '''simple docstring'''




            _UpperCamelCase							       =			parent
            _UpperCamelCase							       =			batch_size
            _UpperCamelCase							       =			is_training
            _UpperCamelCase							       =			use_auxiliary_loss
            _UpperCamelCase							       =			num_queries
            _UpperCamelCase							       =			num_channels
            _UpperCamelCase							       =			min_size
            _UpperCamelCase							       =			max_size
            _UpperCamelCase							       =			num_labels
            _UpperCamelCase							       =			hidden_dim
            _UpperCamelCase							       =			hidden_dim





       def snake_case__		(				self  :						Optional[int]	)       ->	int:

            '''simple docstring'''




            _UpperCamelCase							       =			floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]	).to(
                lowerCAmelCase__	)

            _UpperCamelCase							       =			torch.ones([self.batch_size, self.min_size, self.max_size]			,					device=lowerCAmelCase__	)

            _UpperCamelCase							       =			(
                torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]			,					device=lowerCAmelCase__	) > 0.5
            ).float()
            _UpperCamelCase							       =			(torch.rand((self.batch_size, self.num_labels)			,					device=lowerCAmelCase__	) > 0.5).long()

            _UpperCamelCase							       =			self.get_config()
            return config, pixel_values, pixel_mask, mask_labels, class_labels





       def snake_case__		(				self  :						str	)       ->	str:

            '''simple docstring'''




            _UpperCamelCase							       =			MaskaFormerConfig(
                hidden_size=self.hidden_dim			,					)
            _UpperCamelCase							       =			self.num_queries
            _UpperCamelCase							       =			self.num_labels

            _UpperCamelCase							       =			[1, 1, 1, 1]
            _UpperCamelCase							       =			self.num_channels

            _UpperCamelCase							       =			64
            _UpperCamelCase							       =			128
            _UpperCamelCase							       =			self.hidden_dim
            _UpperCamelCase							       =			self.hidden_dim
            _UpperCamelCase							       =			self.hidden_dim
            return config





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

            '''simple docstring'''




            _UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase							       =			self.prepare_config_and_inputs()
            _UpperCamelCase							       =			{'''pixel_values''': pixel_values, '''pixel_mask''': pixel_mask}
            return config, inputs_dict





       def snake_case__		(				self  :						List[Any]			,					lowerCAmelCase__  :						Any			,					lowerCAmelCase__  :						int	)       ->	Optional[int]:

            '''simple docstring'''




            _UpperCamelCase							       =			output.encoder_hidden_states
            _UpperCamelCase							       =			output.pixel_decoder_hidden_states
            _UpperCamelCase							       =			output.transformer_decoder_hidden_states

            self.parent.assertTrue(len(lowerCAmelCase__	)			,					len(config.backbone_config.depths	)	)
            self.parent.assertTrue(len(lowerCAmelCase__	)			,					len(config.backbone_config.depths	)	)
            self.parent.assertTrue(len(lowerCAmelCase__	)			,					config.decoder_layers	)





       def snake_case__		(				self  :						Optional[Any]			,					lowerCAmelCase__  :						Dict			,					lowerCAmelCase__  :						List[str]			,					lowerCAmelCase__  :						Dict			,					lowerCAmelCase__  :						str=False	)       ->	Tuple:

            '''simple docstring'''




            with torch.no_grad():
                 _UpperCamelCase							       =			MaskaFormerModel(config=lowerCAmelCase__	)
                 model.to(lowerCAmelCase__	)
                 model.eval()

                 _UpperCamelCase							       =			model(pixel_values=lowerCAmelCase__			,					pixel_mask=lowerCAmelCase__	)
                 _UpperCamelCase							       =			model(lowerCAmelCase__			,					output_hidden_states=lowerCAmelCase__	)

            self.parent.assertEqual(
                output.transformer_decoder_last_hidden_state.shape			,					(self.batch_size, self.num_queries, self.hidden_dim)			,					)
            # let's ensure the other two hidden state exists
            self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None	)
            self.parent.assertTrue(output.encoder_last_hidden_state is not None	)

            if output_hidden_states:
                 self.check_output_hidden_state(lowerCAmelCase__			,					lowerCAmelCase__	)







       def snake_case__		(				self  :						Union[str, Any]			,					lowerCAmelCase__  :						str			,					lowerCAmelCase__  :						int			,					lowerCAmelCase__  :						List[str]			,					lowerCAmelCase__  :						Optional[int]			,					lowerCAmelCase__  :						Tuple	)       ->	Tuple:

            '''simple docstring'''




            _UpperCamelCase							       =			MaskaFormerForUniversalSegmentation(config=lowerCAmelCase__	)
            model.to(lowerCAmelCase__	)
            model.eval()

            def comm_check_on_output(lowerCAmelCase__  :						int	):
                 # let's still check that all the required stuff is there
                 self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None	)
                 self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None	)
                 self.parent.assertTrue(result.encoder_last_hidden_state is not None	)
                 # okay, now we need to check the logits shape
                 # due to the encoder compression, masks have a //4 spatial size
                 self.parent.assertEqual(
                     result.masks_queries_logits.shape			,					(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)			,					)
                 # + 1 for null class
                 self.parent.assertEqual(
                     result.class_queries_logits.shape			,					(self.batch_size, self.num_queries, self.num_labels + 1)	)

            with torch.no_grad():
                 _UpperCamelCase							       =			model(pixel_values=lowerCAmelCase__			,					pixel_mask=lowerCAmelCase__	)
                 _UpperCamelCase							       =			model(lowerCAmelCase__	)

                 comm_check_on_output(lowerCAmelCase__	)

                 _UpperCamelCase							       =			model(
                     pixel_values=lowerCAmelCase__			,					pixel_mask=lowerCAmelCase__			,					mask_labels=lowerCAmelCase__			,					class_labels=lowerCAmelCase__	)

            comm_check_on_output(lowerCAmelCase__	)

            self.parent.assertTrue(result.loss is not None	)
            self.parent.assertEqual(result.loss.shape			,					torch.Size([1]	)	)

@require_torch
class 							__lowerCAmelCase	(       __magic_name__		,		__magic_name__		,		unittest.TestCase					):






       """simple docstring"""
       _snake_case      :	Optional[int]									=       (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
       _snake_case      :	List[Any]									=       {'feature-extraction': MaskaFormerModel} if is_torch_available() else {}

       _snake_case      :	Any									=       False
       _snake_case      :	List[str]									=       False
       _snake_case      :	Optional[int]									=       False
       _snake_case      :	int									=       False





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

            '''simple docstring'''




            _UpperCamelCase							       =			MaskaFormerModelTester(self	)
            _UpperCamelCase							       =			ConfigTester(self			,					config_class=lowerCAmelCase__			,					has_text_modality=lowerCAmelCase__	)





       def snake_case__		(				self  :						int	)       ->	List[str]:

            '''simple docstring'''




            self.config_tester.run_common_tests()





       def snake_case__		(				self  :						int	)       ->	List[str]:

            '''simple docstring'''




            _UpperCamelCase       ,							_UpperCamelCase							       =			self.model_tester.prepare_config_and_inputs_for_common()
            self.model_tester.create_and_check_maskaformer_model(lowerCAmelCase__			,					**lowerCAmelCase__			,					output_hidden_states=lowerCAmelCase__	)





       def snake_case__		(				self  :						List[str]	)       ->	str:

            '''simple docstring'''




            _UpperCamelCase							       =			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*lowerCAmelCase__	)





       @unittest.skip(reason='''Mask2Former does not use inputs_embeds'''	)
       def snake_case__		(				self  :						Union[str, Any]	)       ->	List[str]:

            '''simple docstring'''




            pass





       @unittest.skip(reason='''Mask2Former does not have a get_input_embeddings method'''	)
       def snake_case__		(				self  :						Union[str, Any]	)       ->	Tuple:

            '''simple docstring'''




            pass





       @unittest.skip(reason='''Mask2Former is not a generative model'''	)
       def snake_case__		(				self  :						Optional[int]	)       ->	int:

            '''simple docstring'''




            pass





       @unittest.skip(reason='''Mask2Former does not use token embeddings'''	)
       def snake_case__		(				self  :						int	)       ->	List[str]:

            '''simple docstring'''




            pass





       @require_torch_multi_gpu
       @unittest.skip(
           reason='''Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''	)
       def snake_case__		(				self  :						Dict	)       ->	List[Any]:

            '''simple docstring'''




            pass





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

            '''simple docstring'''




            pass





       def snake_case__		(				self  :						List[Any]	)       ->	Any:

            '''simple docstring'''




            _UpperCamelCase       ,							_UpperCamelCase							       =			self.model_tester.prepare_config_and_inputs_for_common()

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

                 _UpperCamelCase							       =			['''pixel_values''']
                 self.assertListEqual(arg_names[:1]			,					lowerCAmelCase__	)





       @slow
       def snake_case__		(				self  :						List[Any]	)       ->	str:

            '''simple docstring'''




            for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
                 _UpperCamelCase							       =			MaskaFormerModel.from_pretrained(lowerCAmelCase__	)
                 self.assertIsNotNone(lowerCAmelCase__	)





       def snake_case__		(				self  :						int	)       ->	Tuple:

            '''simple docstring'''




            _UpperCamelCase							       =			(self.model_tester.min_size,) * 2
            _UpperCamelCase							       =			{
                '''pixel_values''': torch.randn((2, 3, *size)			,					device=lowerCAmelCase__	),
                '''mask_labels''': torch.randn((2, 10, *size)			,					device=lowerCAmelCase__	),
                '''class_labels''': torch.zeros(2			,					10			,					device=lowerCAmelCase__	).long(),
            }
            _UpperCamelCase							       =			self.model_tester.get_config()

            _UpperCamelCase							       =			MaskaFormerForUniversalSegmentation(lowerCAmelCase__	).to(lowerCAmelCase__	)
            _UpperCamelCase							       =			model(**lowerCAmelCase__	)
            self.assertTrue(outputs.loss is not None	)





       def snake_case__		(				self  :						Tuple	)       ->	Optional[Any]:

            '''simple docstring'''




            _UpperCamelCase       ,							_UpperCamelCase							       =			self.model_tester.prepare_config_and_inputs_for_common()
            self.model_tester.create_and_check_maskaformer_model(lowerCAmelCase__			,					**lowerCAmelCase__			,					output_hidden_states=lowerCAmelCase__	)





       def snake_case__		(				self  :						str	)       ->	Optional[int]:

            '''simple docstring'''




            _UpperCamelCase       ,							_UpperCamelCase							       =			self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                 _UpperCamelCase							       =			model_class(lowerCAmelCase__	).to(lowerCAmelCase__	)
                 _UpperCamelCase							       =			model(**lowerCAmelCase__			,					output_attentions=lowerCAmelCase__	)
                 self.assertTrue(outputs.attentions is not None	)





       def snake_case__		(				self  :						Tuple	)       ->	Dict:

            '''simple docstring'''




            if not self.model_tester.is_training:
                 return

            _UpperCamelCase							       =			self.all_model_classes[1]
            _UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase							       =			self.model_tester.prepare_config_and_inputs()

            _UpperCamelCase							       =			model_class(lowerCAmelCase__	)
            model.to(lowerCAmelCase__	)
            model.train()

            _UpperCamelCase							       =			model(lowerCAmelCase__			,					mask_labels=lowerCAmelCase__			,					class_labels=lowerCAmelCase__	).loss
            loss.backward()







       def snake_case__		(				self  :						Optional[Any]	)       ->	Optional[Any]:

            '''simple docstring'''




            _UpperCamelCase							       =			self.all_model_classes[1]
            _UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase       ,							_UpperCamelCase							       =			self.model_tester.prepare_config_and_inputs()
            _UpperCamelCase							       =			True
            _UpperCamelCase							       =			True

            _UpperCamelCase							       =			model_class(lowerCAmelCase__	).to(lowerCAmelCase__	)
            model.train()

            _UpperCamelCase							       =			model(lowerCAmelCase__			,					mask_labels=lowerCAmelCase__			,					class_labels=lowerCAmelCase__	)

            _UpperCamelCase							       =			outputs.encoder_hidden_states[0]
            encoder_hidden_states.retain_grad()

            _UpperCamelCase							       =			outputs.pixel_decoder_hidden_states[0]
            pixel_decoder_hidden_states.retain_grad()

            _UpperCamelCase							       =			outputs.transformer_decoder_hidden_states[0]
            transformer_decoder_hidden_states.retain_grad()

            _UpperCamelCase							       =			outputs.attentions[0]
            attentions.retain_grad()

            outputs.loss.backward(retain_graph=lowerCAmelCase__	)

            self.assertIsNotNone(encoder_hidden_states.grad	)
            self.assertIsNotNone(pixel_decoder_hidden_states.grad	)
            self.assertIsNotNone(transformer_decoder_hidden_states.grad	)
            self.assertIsNotNone(attentions.grad	)

lowercase__		:     Tuple			=							1E-4







def    a__							(							)						->						Any:





     """simple docstring"""

     _UpperCamelCase							       =			Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''' )
     return image




@require_vision
@slow
class 							__lowerCAmelCase	(       unittest.TestCase					):






       """simple docstring"""





       @cached_property
       def snake_case__		(				self  :						List[str]	)       ->	Any:

            '''simple docstring'''




            return "facebook/mask2former-swin-small-coco-instance"





       @cached_property
       def snake_case__		(				self  :						List[str]	)       ->	List[str]:

            '''simple docstring'''




            return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints	) if is_vision_available() else None





       def snake_case__		(				self  :						List[str]	)       ->	Tuple:

            '''simple docstring'''




            _UpperCamelCase							       =			MaskaFormerModel.from_pretrained(self.model_checkpoints	).to(lowerCAmelCase__	)
            _UpperCamelCase							       =			self.default_image_processor
            _UpperCamelCase							       =			prepare_img()
            _UpperCamelCase							       =			image_processor(lowerCAmelCase__			,					return_tensors='''pt'''	).to(lowerCAmelCase__	)
            _UpperCamelCase							       =			inputs['''pixel_values'''].shape
            # check size is divisible by 32
            self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0	)
            # check size
            self.assertEqual(lowerCAmelCase__			,					(1, 3, 384, 384)	)

            with torch.no_grad():
                 _UpperCamelCase							       =			model(**lowerCAmelCase__	)

            _UpperCamelCase							       =			torch.tensor(
                [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]	).to(lowerCAmelCase__	)
            self.assertTrue(
                torch.allclose(
                    outputs.encoder_last_hidden_state[0, 0, :3, :3]			,					lowerCAmelCase__			,					atol=lowerCAmelCase__	)	)

            _UpperCamelCase							       =			torch.tensor(
                [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]	).to(lowerCAmelCase__	)
            self.assertTrue(
                torch.allclose(
                    outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]			,					lowerCAmelCase__			,					atol=lowerCAmelCase__	)	)

            _UpperCamelCase							       =			torch.tensor(
                [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]	).to(lowerCAmelCase__	)
            self.assertTrue(
                torch.allclose(
                    outputs.transformer_decoder_last_hidden_state[0, :3, :3]			,					lowerCAmelCase__			,					atol=lowerCAmelCase__	)	)





       def snake_case__		(				self  :						Tuple	)       ->	List[Any]:

            '''simple docstring'''




            _UpperCamelCase							       =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints	).to(lowerCAmelCase__	).eval()
            _UpperCamelCase							       =			self.default_image_processor
            _UpperCamelCase							       =			prepare_img()
            _UpperCamelCase							       =			image_processor(lowerCAmelCase__			,					return_tensors='''pt'''	).to(lowerCAmelCase__	)
            _UpperCamelCase							       =			inputs['''pixel_values'''].shape
            # check size is divisible by 32
            self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0	)
            # check size
            self.assertEqual(lowerCAmelCase__			,					(1, 3, 384, 384)	)

            with torch.no_grad():
                 _UpperCamelCase							       =			model(**lowerCAmelCase__	)
            # masks_queries_logits
            _UpperCamelCase							       =			outputs.masks_queries_logits
            self.assertEqual(
                masks_queries_logits.shape			,					(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)	)
            _UpperCamelCase							       =			[
                [-8.7839, -9.0056, -8.8121],
                [-7.4104, -7.0313, -6.5401],
                [-6.6105, -6.3427, -6.4675],
            ]
            _UpperCamelCase							       =			torch.tensor(lowerCAmelCase__	).to(lowerCAmelCase__	)
            self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]			,					lowerCAmelCase__			,					atol=lowerCAmelCase__	)	)
            # class_queries_logits
            _UpperCamelCase							       =			outputs.class_queries_logits
            self.assertEqual(class_queries_logits.shape			,					(1, model.config.num_queries, model.config.num_labels + 1)	)
            _UpperCamelCase							       =			torch.tensor(
                [
                    [1.8324, -8.0835, -4.1922],
                    [0.8450, -9.0050, -3.6053],
                    [0.3045, -7.7293, -3.0275],
                ]	).to(lowerCAmelCase__	)
            self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]			,					lowerCAmelCase__			,					atol=lowerCAmelCase__	)	)







       def snake_case__		(				self  :						Tuple	)       ->	int:

            '''simple docstring'''




            _UpperCamelCase							       =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints	).to(lowerCAmelCase__	).eval()
            _UpperCamelCase							       =			self.default_image_processor

            _UpperCamelCase							       =			image_processor(
                [np.zeros((3, 800, 1333)	), np.zeros((3, 800, 1333)	)]			,					segmentation_maps=[np.zeros((384, 384)	).astype(np.floataa	), np.zeros((384, 384)	).astype(np.floataa	)]			,					return_tensors='''pt'''			,					)

            _UpperCamelCase							       =			inputs['''pixel_values'''].to(lowerCAmelCase__	)
            _UpperCamelCase							       =			[el.to(lowerCAmelCase__	) for el in inputs['''mask_labels''']]
            _UpperCamelCase							       =			[el.to(lowerCAmelCase__	) for el in inputs['''class_labels''']]

            with torch.no_grad():
                 _UpperCamelCase							       =			model(**lowerCAmelCase__	)

            self.assertTrue(outputs.loss is not None	)



