





'''simple docstring'''






import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import DetrConfig, MaskFormerConfig, SwinConfig, 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 MaskFormerForInstanceSegmentation, MaskFormerModel

      if is_vision_available():
            from transformers import MaskFormerImageProcessor

if is_vision_available():
      from PIL import Image




class  UpperCAmelCase     :




      def __init__(			self			: Optional[int]			,			__snake_case			: int			,			__snake_case			: Optional[Any]=2			,			__snake_case			: int=True			,			__snake_case			: str=False			,			__snake_case			: List[str]=10			,			__snake_case			: Union[str, Any]=3			,			__snake_case			: List[Any]=32 * 4			,			__snake_case			: str=32 * 6			,			__snake_case			: int=4			,			__snake_case			: str=32			,			)				->      str:
            _lowerCAmelCase							     =							parent
            _lowerCAmelCase							     =							batch_size
            _lowerCAmelCase							     =							is_training
            _lowerCAmelCase							     =							use_auxiliary_loss
            _lowerCAmelCase							     =							num_queries
            _lowerCAmelCase							     =							num_channels
            _lowerCAmelCase							     =							min_size
            _lowerCAmelCase							     =							max_size
            _lowerCAmelCase							     =							num_labels
            _lowerCAmelCase							     =							mask_feature_size




      def 				lowercase__    (			self			: Optional[int]  )				->      Optional[int]:
            _lowerCAmelCase							     =							floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]  ).to(
                __snake_case  )

            _lowerCAmelCase							     =							torch.ones([self.batch_size, self.min_size, self.max_size]			,			device=__snake_case  )

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

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




      def 				lowercase__    (			self			: Any  )				->      Union[str, Any]:
            return MaskFormerConfig.from_backbone_and_decoder_configs(
                backbone_config=SwinConfig(
                    depths=[1, 1, 1, 1]			,			)			,			decoder_config=DetrConfig(
                    decoder_ffn_dim=1_28			,			num_queries=self.num_queries			,			decoder_attention_heads=2			,			d_model=self.mask_feature_size			,			)			,			mask_feature_size=self.mask_feature_size			,			fpn_feature_size=self.mask_feature_size			,			num_channels=self.num_channels			,			num_labels=self.num_labels			,			)




      def 				lowercase__    (			self			: Tuple  )				->      Tuple:
            _lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase							     =							self.prepare_config_and_inputs()
            _lowerCAmelCase							     =							{"""pixel_values""": pixel_values, """pixel_mask""": pixel_mask}
            return config, inputs_dict




      def 				lowercase__    (			self			: List[Any]			,			__snake_case			: str			,			__snake_case			: Optional[int]  )				->      List[Any]:
            _lowerCAmelCase							     =							output.encoder_hidden_states
            _lowerCAmelCase							     =							output.pixel_decoder_hidden_states
            _lowerCAmelCase							     =							output.transformer_decoder_hidden_states

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




      def 				lowercase__    (			self			: str			,			__snake_case			: Any			,			__snake_case			: Union[str, Any]			,			__snake_case			: Optional[int]			,			__snake_case			: Dict=False  )				->      Dict:
            with torch.no_grad():
                  _lowerCAmelCase							     =							MaskFormerModel(config=__snake_case  )
                  model.to(__snake_case  )
                  model.eval()

                  _lowerCAmelCase							     =							model(pixel_values=__snake_case			,			pixel_mask=__snake_case  )
                  _lowerCAmelCase							     =							model(__snake_case			,			output_hidden_states=__snake_case  )
            # the correct shape of output.transformer_decoder_hidden_states ensure the correcteness of the
            # encoder and pixel decoder
            self.parent.assertEqual(
                output.transformer_decoder_last_hidden_state.shape			,			(self.batch_size, self.num_queries, self.mask_feature_size)			,			)
            # 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(__snake_case			,			__snake_case  )




      def 				lowercase__    (			self			: Union[str, Any]			,			__snake_case			: int			,			__snake_case			: int			,			__snake_case			: Dict			,			__snake_case			: Dict			,			__snake_case			: str  )				->      str:
            _lowerCAmelCase							     =							MaskFormerForInstanceSegmentation(config=__snake_case  )
            model.to(__snake_case  )
            model.eval()

            def comm_check_on_output(__snake_case			: List[str]  ):
                  # 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():
                  _lowerCAmelCase							     =							model(pixel_values=__snake_case			,			pixel_mask=__snake_case  )
                  _lowerCAmelCase							     =							model(__snake_case  )

                  comm_check_on_output(__snake_case  )

                  _lowerCAmelCase							     =							model(
                      pixel_values=__snake_case			,			pixel_mask=__snake_case			,			mask_labels=__snake_case			,			class_labels=__snake_case  )

            comm_check_on_output(__snake_case  )

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




@require_torch
class  UpperCAmelCase     (			snake_case_							,		snake_case_							,		unittest.TestCase   ):
      _lowercase:			Any						     =   (MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
      _lowercase:			int						     =   (
          {'''feature-extraction''': MaskFormerModel, '''image-segmentation''': MaskFormerForInstanceSegmentation}
          if is_torch_available()
          else {}
      )

      _lowercase:			Optional[int]						     =   False
      _lowercase:			Union[str, Any]						     =   False
      _lowercase:			Dict						     =   False
      _lowercase:			Union[str, Any]						     =   False




      def 				lowercase__    (			self			: Tuple  )				->      List[Any]:
            _lowerCAmelCase							     =							MaskFormerModelTester(self  )
            _lowerCAmelCase							     =							ConfigTester(self			,			config_class=__snake_case			,			has_text_modality=__snake_case  )




      def 				lowercase__    (			self			: Dict  )				->      Dict:
            self.config_tester.run_common_tests()




      def 				lowercase__    (			self			: Optional[Any]  )				->      str:
            _lowerCAmelCase     ,							_lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs_for_common()
            self.model_tester.create_and_check_maskformer_model(__snake_case			,			**__snake_case			,			output_hidden_states=__snake_case  )




      def 				lowercase__    (			self			: str  )				->      Optional[int]:
            _lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_maskformer_instance_segmentation_head_model(*__snake_case  )




      @unittest.skip(reason="""MaskFormer does not use inputs_embeds"""  )
      def 				lowercase__    (			self			: List[Any]  )				->      List[str]:
            pass




      @unittest.skip(reason="""MaskFormer does not have a get_input_embeddings method"""  )
      def 				lowercase__    (			self			: Optional[Any]  )				->      Optional[int]:
            pass




      @unittest.skip(reason="""MaskFormer is not a generative model"""  )
      def 				lowercase__    (			self			: Union[str, Any]  )				->      List[str]:
            pass




      @unittest.skip(reason="""MaskFormer does not use token embeddings"""  )
      def 				lowercase__    (			self			: List[str]  )				->      Any:
            pass




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




      @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""  )
      def 				lowercase__    (			self			: Optional[Any]  )				->      Dict:
            pass




      def 				lowercase__    (			self			: Tuple  )				->      Union[str, Any]:
            _lowerCAmelCase     ,							_lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs_for_common()

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

                  _lowerCAmelCase							     =							["""pixel_values"""]
                  self.assertListEqual(arg_names[:1]			,			__snake_case  )




      @slow
      def 				lowercase__    (			self			: Optional[Any]  )				->      str:
            for model_name in ["facebook/maskformer-swin-small-coco"]:
                  _lowerCAmelCase							     =							MaskFormerModel.from_pretrained(__snake_case  )
                  self.assertIsNotNone(__snake_case  )




      def 				lowercase__    (			self			: str  )				->      int:
            _lowerCAmelCase							     =							(self.model_tester.min_size,) * 2
            _lowerCAmelCase							     =							{
                """pixel_values""": torch.randn((2, 3, *size)			,			device=__snake_case  ),
                """mask_labels""": torch.randn((2, 10, *size)			,			device=__snake_case  ),
                """class_labels""": torch.zeros(2			,			10			,			device=__snake_case  ).long(),
            }

            _lowerCAmelCase							     =							MaskFormerForInstanceSegmentation(MaskFormerConfig()  ).to(__snake_case  )
            _lowerCAmelCase							     =							model(**__snake_case  )
            self.assertTrue(outputs.loss is not None  )




      def 				lowercase__    (			self			: Optional[Any]  )				->      List[Any]:
            _lowerCAmelCase     ,							_lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs_for_common()
            self.model_tester.create_and_check_maskformer_model(__snake_case			,			**__snake_case			,			output_hidden_states=__snake_case  )




      def 				lowercase__    (			self			: str  )				->      Optional[int]:
            _lowerCAmelCase     ,							_lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs_for_common()

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




      def 				lowercase__    (			self			: Tuple  )				->      str:
            if not self.model_tester.is_training:
                  return
            # only MaskFormerForInstanceSegmentation has the loss
            _lowerCAmelCase							     =							self.all_model_classes[1]
            _lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs()

            _lowerCAmelCase							     =							model_class(__snake_case  )
            model.to(__snake_case  )
            model.train()

            _lowerCAmelCase							     =							model(__snake_case			,			mask_labels=__snake_case			,			class_labels=__snake_case  ).loss
            loss.backward()




      def 				lowercase__    (			self			: Dict  )				->      int:
            # only MaskFormerForInstanceSegmentation has the loss
            _lowerCAmelCase							     =							self.all_model_classes[1]
            _lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase     ,							_lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs()
            _lowerCAmelCase							     =							True
            _lowerCAmelCase							     =							True

            _lowerCAmelCase							     =							model_class(__snake_case  )
            model.to(__snake_case  )
            model.train()

            _lowerCAmelCase							     =							model(__snake_case			,			mask_labels=__snake_case			,			class_labels=__snake_case  )

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

            _lowerCAmelCase							     =							outputs.pixel_decoder_hidden_states[0]
            pixel_decoder_hidden_states.retain_grad()
            # we requires_grad=True in inputs_embeds (line 2152), the original implementation don't
            _lowerCAmelCase							     =							outputs.transformer_decoder_hidden_states[0]
            transformer_decoder_hidden_states.retain_grad()

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

            outputs.loss.backward(retain_graph=__snake_case  )

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




A__				:   int      =1e-4
def    UpperCamelCase__      (		):





      """simple docstring"""


      _lowerCAmelCase							     =							Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""  )
      return image





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




      @cached_property
      def 				lowercase__    (			self			: Dict  )				->      Dict:
            return (
                MaskFormerImageProcessor.from_pretrained("""facebook/maskformer-swin-small-coco"""  )
                if is_vision_available()
                else None
            )




      def 				lowercase__    (			self			: str  )				->      Union[str, Any]:
            _lowerCAmelCase							     =							MaskFormerModel.from_pretrained("""facebook/maskformer-swin-small-coco"""  ).to(__snake_case  )
            _lowerCAmelCase							     =							self.default_image_processor
            _lowerCAmelCase							     =							prepare_img()
            _lowerCAmelCase							     =							image_processor(__snake_case			,			return_tensors="""pt"""  ).to(__snake_case  )
            _lowerCAmelCase							     =							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(__snake_case			,			(1, 3, 8_00, 10_88)  )

            with torch.no_grad():
                  _lowerCAmelCase							     =							model(**__snake_case  )

            _lowerCAmelCase							     =							torch.tensor(
                [[-0.04_82, 0.92_28, 0.49_51], [-0.25_47, 0.80_17, 0.85_27], [-0.00_69, 0.33_85, -0.00_89]]  ).to(__snake_case  )
            self.assertTrue(
                torch.allclose(
                    outputs.encoder_last_hidden_state[0, 0, :3, :3]			,			__snake_case			,			atol=__snake_case  )  )

            _lowerCAmelCase							     =							torch.tensor(
                [[-0.84_22, -0.84_34, -0.97_18], [-1.01_44, -0.55_65, -0.41_95], [-1.00_38, -0.44_84, -0.19_61]]  ).to(__snake_case  )
            self.assertTrue(
                torch.allclose(
                    outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]			,			__snake_case			,			atol=__snake_case  )  )

            _lowerCAmelCase							     =							torch.tensor(
                [[0.28_52, -0.01_59, 0.97_35], [0.62_54, 0.18_58, 0.85_29], [-0.06_80, -0.41_16, 1.84_13]]  ).to(__snake_case  )
            self.assertTrue(
                torch.allclose(
                    outputs.transformer_decoder_last_hidden_state[0, :3, :3]			,			__snake_case			,			atol=__snake_case  )  )




      def 				lowercase__    (			self			: Any  )				->      Tuple:
            _lowerCAmelCase							     =							(
                MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-swin-small-coco"""  )
                .to(__snake_case  )
                .eval()
            )
            _lowerCAmelCase							     =							self.default_image_processor
            _lowerCAmelCase							     =							prepare_img()
            _lowerCAmelCase							     =							image_processor(__snake_case			,			return_tensors="""pt"""  ).to(__snake_case  )
            _lowerCAmelCase							     =							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(__snake_case			,			(1, 3, 8_00, 10_88)  )

            with torch.no_grad():
                  _lowerCAmelCase							     =							model(**__snake_case  )
            # masks_queries_logits
            _lowerCAmelCase							     =							outputs.masks_queries_logits
            self.assertEqual(
                masks_queries_logits.shape			,			(1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)			,			)
            _lowerCAmelCase							     =							[
                [-1.3_73_71_24, -1.7_72_49_37, -1.9_36_42_33],
                [-1.5_97_72_81, -1.9_86_79_39, -2.1_52_36_95],
                [-1.5_79_53_98, -1.9_26_98_32, -2.09_39_42],
            ]
            _lowerCAmelCase							     =							torch.tensor(__snake_case  ).to(__snake_case  )
            self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]			,			__snake_case			,			atol=__snake_case  )  )
            # class_queries_logits
            _lowerCAmelCase							     =							outputs.class_queries_logits
            self.assertEqual(
                class_queries_logits.shape			,			(1, model.config.decoder_config.num_queries, model.config.num_labels + 1)  )
            _lowerCAmelCase							     =							torch.tensor(
                [
                    [1.6_5_1_2E0_0, -5.2_5_7_2E0_0, -3.3_5_1_9E0_0],
                    [3.6_1_6_9E-0_2, -5.9_0_2_5E0_0, -2.9_3_1_3E0_0],
                    [1.0_7_6_6E-0_4, -7.7_6_3_0E0_0, -5.1_2_6_3E0_0],
                ]  ).to(__snake_case  )
            self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]			,			__snake_case			,			atol=__snake_case  )  )




      def 				lowercase__    (			self			: Tuple  )				->      Optional[Any]:
            _lowerCAmelCase							     =							(
                MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-resnet101-coco-stuff"""  )
                .to(__snake_case  )
                .eval()
            )
            _lowerCAmelCase							     =							self.default_image_processor
            _lowerCAmelCase							     =							prepare_img()
            _lowerCAmelCase							     =							image_processor(__snake_case			,			return_tensors="""pt"""  ).to(__snake_case  )
            _lowerCAmelCase							     =							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(__snake_case			,			(1, 3, 8_00, 10_88)  )

            with torch.no_grad():
                  _lowerCAmelCase							     =							model(**__snake_case  )
            # masks_queries_logits
            _lowerCAmelCase							     =							outputs.masks_queries_logits
            self.assertEqual(
                masks_queries_logits.shape			,			(1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)			,			)
            _lowerCAmelCase							     =							[[-0.90_46, -2.63_66, -4.60_62], [-3.41_79, -5.78_90, -8.80_57], [-4.91_79, -7.65_60, -10.77_11]]
            _lowerCAmelCase							     =							torch.tensor(__snake_case  ).to(__snake_case  )
            self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]			,			__snake_case			,			atol=__snake_case  )  )
            # class_queries_logits
            _lowerCAmelCase							     =							outputs.class_queries_logits
            self.assertEqual(
                class_queries_logits.shape			,			(1, model.config.decoder_config.num_queries, model.config.num_labels + 1)  )
            _lowerCAmelCase							     =							torch.tensor(
                [[4.71_88, -3.25_85, -2.88_57], [6.68_71, -2.91_81, -1.24_87], [7.24_49, -2.27_64, -2.18_74]]  ).to(__snake_case  )
            self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]			,			__snake_case			,			atol=__snake_case  )  )




      def 				lowercase__    (			self			: Optional[int]  )				->      Union[str, Any]:
            _lowerCAmelCase							     =							(
                MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-swin-small-coco"""  )
                .to(__snake_case  )
                .eval()
            )
            _lowerCAmelCase							     =							self.default_image_processor

            _lowerCAmelCase							     =							image_processor(
                [np.zeros((3, 8_00, 13_33)  ), np.zeros((3, 8_00, 13_33)  )]			,			segmentation_maps=[np.zeros((3_84, 3_84)  ).astype(np.floataa  ), np.zeros((3_84, 3_84)  ).astype(np.floataa  )]			,			return_tensors="""pt"""			,			)

            _lowerCAmelCase							     =							inputs["""pixel_values"""].to(__snake_case  )
            _lowerCAmelCase							     =							[el.to(__snake_case  ) for el in inputs["""mask_labels"""]]
            _lowerCAmelCase							     =							[el.to(__snake_case  ) for el in inputs["""class_labels"""]]

            with torch.no_grad():
                  _lowerCAmelCase							     =							model(**__snake_case  )

            self.assertTrue(outputs.loss is not None  )








'''simple docstring'''




import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD





class  A				(							nn.Module						):
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		0.0
     __magic_name__						=		1
     __magic_name__						=		1
     __magic_name__						=		True
     __magic_name__						=		False
     __magic_name__						=		False
     __magic_name__						=		False
     __magic_name__						=		jnp.floataa





     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         A       :   Union[str, Any]					=						[]
         A       :   Union[str, Any]					=						[]

         for i in range(self.num_layers	):
             A       :   Any					=						self.in_channels if i == 0 else self.out_channels

             A       :   Optional[Any]					=						FlaxResnetBlockaD(
                 in_channels=SCREAMING_SNAKE_CASE						,			out_channels=self.out_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
             resnets.append(SCREAMING_SNAKE_CASE	)

             A       :   Optional[int]					=						FlaxTransformeraDModel(
                 in_channels=self.out_channels						,			n_heads=self.num_attention_heads						,			d_head=self.out_channels // self.num_attention_heads						,			depth=1						,			use_linear_projection=self.use_linear_projection						,			only_cross_attention=self.only_cross_attention						,			use_memory_efficient_attention=self.use_memory_efficient_attention						,			dtype=self.dtype						,			)
             attentions.append(SCREAMING_SNAKE_CASE	)

         A       :   Union[str, Any]					=						resnets
         A       :   Union[str, Any]					=						attentions

         if self.add_downsample:
             A       :   int					=						FlaxDownsampleaD(self.out_channels						,			dtype=self.dtype	)





     def __call__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=True	) ->				Union[str, Any]:






         """simple docstring"""






         A       :   Optional[Any]					=						()

         for resnet, attn in zip(self.resnets						,			self.attentions	):
             A       :   int					=						resnet(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)
             A       :   Dict					=						attn(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)
             output_states += (hidden_states,)

         if self.add_downsample:
             A       :   Optional[Any]					=						self.downsamplers_a(SCREAMING_SNAKE_CASE	)
             output_states += (hidden_states,)

         return hidden_states, output_states





class  A				(							nn.Module						):
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		0.0
     __magic_name__						=		1
     __magic_name__						=		True
     __magic_name__						=		jnp.floataa





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






         """simple docstring"""






         A       :   Optional[Any]					=						[]

         for i in range(self.num_layers	):
             A       :   Optional[Any]					=						self.in_channels if i == 0 else self.out_channels

             A       :   List[str]					=						FlaxResnetBlockaD(
                 in_channels=SCREAMING_SNAKE_CASE						,			out_channels=self.out_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
             resnets.append(SCREAMING_SNAKE_CASE	)
         A       :   Dict					=						resnets

         if self.add_downsample:
             A       :   Dict					=						FlaxDownsampleaD(self.out_channels						,			dtype=self.dtype	)





     def __call__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=True	) ->				Optional[Any]:






         """simple docstring"""






         A       :   str					=						()

         for resnet in self.resnets:
             A       :   Optional[int]					=						resnet(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)
             output_states += (hidden_states,)

         if self.add_downsample:
             A       :   Optional[int]					=						self.downsamplers_a(SCREAMING_SNAKE_CASE	)
             output_states += (hidden_states,)

         return hidden_states, output_states





class  A				(							nn.Module						):
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		0.0
     __magic_name__						=		1
     __magic_name__						=		1
     __magic_name__						=		True
     __magic_name__						=		False
     __magic_name__						=		False
     __magic_name__						=		False
     __magic_name__						=		jnp.floataa





     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         A       :   Optional[Any]					=						[]
         A       :   Optional[int]					=						[]

         for i in range(self.num_layers	):
             A       :   str					=						self.in_channels if (i == self.num_layers - 1) else self.out_channels
             A       :   Dict					=						self.prev_output_channel if i == 0 else self.out_channels

             A       :   List[str]					=						FlaxResnetBlockaD(
                 in_channels=resnet_in_channels + res_skip_channels						,			out_channels=self.out_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
             resnets.append(SCREAMING_SNAKE_CASE	)

             A       :   int					=						FlaxTransformeraDModel(
                 in_channels=self.out_channels						,			n_heads=self.num_attention_heads						,			d_head=self.out_channels // self.num_attention_heads						,			depth=1						,			use_linear_projection=self.use_linear_projection						,			only_cross_attention=self.only_cross_attention						,			use_memory_efficient_attention=self.use_memory_efficient_attention						,			dtype=self.dtype						,			)
             attentions.append(SCREAMING_SNAKE_CASE	)

         A       :   Dict					=						resnets
         A       :   Optional[Any]					=						attentions

         if self.add_upsample:
             A       :   Optional[int]					=						FlaxUpsampleaD(self.out_channels						,			dtype=self.dtype	)





     def __call__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=True	) ->				Optional[int]:






         """simple docstring"""






         for resnet, attn in zip(self.resnets						,			self.attentions	):
             # pop res hidden states
             A       :   List[str]					=						res_hidden_states_tuple[-1]
             A       :   int					=						res_hidden_states_tuple[:-1]
             A       :   List[str]					=						jnp.concatenate((hidden_states, res_hidden_states)						,			axis=-1	)

             A       :   Union[str, Any]					=						resnet(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)
             A       :   Tuple					=						attn(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)

         if self.add_upsample:
             A       :   Dict					=						self.upsamplers_a(SCREAMING_SNAKE_CASE	)

         return hidden_states





class  A				(							nn.Module						):
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		0.0
     __magic_name__						=		1
     __magic_name__						=		True
     __magic_name__						=		jnp.floataa





     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   int					=						[]

         for i in range(self.num_layers	):
             A       :   List[Any]					=						self.in_channels if (i == self.num_layers - 1) else self.out_channels
             A       :   List[str]					=						self.prev_output_channel if i == 0 else self.out_channels

             A       :   str					=						FlaxResnetBlockaD(
                 in_channels=resnet_in_channels + res_skip_channels						,			out_channels=self.out_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
             resnets.append(SCREAMING_SNAKE_CASE	)

         A       :   List[Any]					=						resnets

         if self.add_upsample:
             A       :   Optional[Any]					=						FlaxUpsampleaD(self.out_channels						,			dtype=self.dtype	)





     def __call__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=True	) ->				Tuple:






         """simple docstring"""






         for resnet in self.resnets:
             # pop res hidden states
             A       :   Optional[int]					=						res_hidden_states_tuple[-1]
             A       :   Optional[Any]					=						res_hidden_states_tuple[:-1]
             A       :   List[Any]					=						jnp.concatenate((hidden_states, res_hidden_states)						,			axis=-1	)

             A       :   Optional[Any]					=						resnet(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)

         if self.add_upsample:
             A       :   List[str]					=						self.upsamplers_a(SCREAMING_SNAKE_CASE	)

         return hidden_states





class  A				(							nn.Module						):
     __magic_name__						=		42
     __magic_name__						=		0.0
     __magic_name__						=		1
     __magic_name__						=		1
     __magic_name__						=		False
     __magic_name__						=		False
     __magic_name__						=		jnp.floataa





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






         """simple docstring"""






         A       :   str					=						[
             FlaxResnetBlockaD(
                 in_channels=self.in_channels						,			out_channels=self.in_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
         ]

         A       :   List[Any]					=						[]

         for _ in range(self.num_layers	):
             A       :   int					=						FlaxTransformeraDModel(
                 in_channels=self.in_channels						,			n_heads=self.num_attention_heads						,			d_head=self.in_channels // self.num_attention_heads						,			depth=1						,			use_linear_projection=self.use_linear_projection						,			use_memory_efficient_attention=self.use_memory_efficient_attention						,			dtype=self.dtype						,			)
             attentions.append(SCREAMING_SNAKE_CASE	)

             A       :   Union[str, Any]					=						FlaxResnetBlockaD(
                 in_channels=self.in_channels						,			out_channels=self.in_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
             resnets.append(SCREAMING_SNAKE_CASE	)

         A       :   List[str]					=						resnets
         A       :   List[str]					=						attentions





     def __call__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=True	) ->				Dict:






         """simple docstring"""






         A       :   Optional[Any]					=						self.resnets[0](SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	)
         for attn, resnet in zip(self.attentions						,			self.resnets[1:]	):
             A       :   Optional[int]					=						attn(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)
             A       :   Union[str, Any]					=						resnet(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)

         return hidden_states




