





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    __snake_case	(							UpperCAmelCase_      ):





       '''simple docstring'''







       lowerCAmelCase__						=					42







@flax_register_to_config
class    __snake_case	(							nn.Module       ,      UpperCAmelCase_       ,      UpperCAmelCase_      ):





       '''simple docstring'''







       lowerCAmelCase__						=					32
       lowerCAmelCase__						=					4
       lowerCAmelCase__						=					4
       lowerCAmelCase__						=					(
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "CrossAttnDownBlock2D",
           "DownBlock2D",
       )
       lowerCAmelCase__						=					("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
       lowerCAmelCase__						=					False
       lowerCAmelCase__						=					(3_20, 6_40, 12_80, 12_80)
       lowerCAmelCase__						=					2
       lowerCAmelCase__						=					8
       lowerCAmelCase__						=					None
       lowerCAmelCase__						=					12_80
       lowerCAmelCase__						=					0.0
       lowerCAmelCase__						=					False
       lowerCAmelCase__						=					jnp.floataa
       lowerCAmelCase__						=					True
       lowerCAmelCase__						=					0
       lowerCAmelCase__						=					False




       def 							UpperCAmelCase__			(  self			:   List[Any]							,    A			:   int		):
         # init input tensors
         __snake_case:						int			  =						(1, self.in_channels, self.sample_size, self.sample_size)
         __snake_case:						Tuple			  =						jnp.zeros(__lowercase							,    dtype=jnp.floataa		)
         __snake_case:						int			  =						jnp.ones((1,)							,    dtype=jnp.intaa		)
         __snake_case:						Tuple			  =						jnp.zeros((1, 1, self.cross_attention_dim)							,    dtype=jnp.floataa		)

         __snake_case:						Optional[Any]			  =						jax.random.split(__lowercase		)
         __snake_case:						List[str]			  =						{'''params''': params_rng, '''dropout''': dropout_rng}

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




       def 							UpperCAmelCase__			(  self			:   str		):
         __snake_case:						List[str]			  =						self.block_out_channels
         __snake_case:						Dict			  =						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:						str			  =						self.num_attention_heads or self.attention_head_dim

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

         # time
         __snake_case:						Union[str, Any]			  =						FlaxTimesteps(
             block_out_channels[0]							,    flip_sin_to_cos=self.flip_sin_to_cos							,    freq_shift=self.config.freq_shift		)
         __snake_case:						List[str]			  =						FlaxTimestepEmbedding(__lowercase							,    dtype=self.dtype		)

         __snake_case:						str			  =						self.only_cross_attention
         if isinstance(__lowercase							,    __lowercase		):
           __snake_case:						Tuple			  =						(only_cross_attention,) * len(self.down_block_types		)

         if isinstance(__lowercase							,    __lowercase		):
           __snake_case:						Any			  =						(num_attention_heads,) * len(self.down_block_types		)

         # down
         __snake_case:						List[Any]			  =						[]
         __snake_case:						Dict			  =						block_out_channels[0]
         for i, down_block_type in enumerate(self.down_block_types		):
           __snake_case:						Dict			  =						output_channel
           __snake_case:						int			  =						block_out_channels[i]
           __snake_case:						List[Any]			  =						i == len(__lowercase		) - 1

           if down_block_type == "CrossAttnDownBlock2D":
             __snake_case:						Tuple			  =						FlaxCrossAttnDownBlockaD(
                 in_channels=__lowercase							,    out_channels=__lowercase							,    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:						List[str]			  =						FlaxDownBlockaD(
                 in_channels=__lowercase							,    out_channels=__lowercase							,    dropout=self.dropout							,    num_layers=self.layers_per_block							,    add_downsample=not is_final_block							,    dtype=self.dtype							,    )

           down_blocks.append(__lowercase		)
         __snake_case:						List[str]			  =						down_blocks

         # mid
         __snake_case:						str			  =						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:						Tuple			  =						[]
         __snake_case:						List[Any]			  =						list(reversed(__lowercase		)		)
         __snake_case:						int			  =						list(reversed(__lowercase		)		)
         __snake_case:						Tuple			  =						list(reversed(__lowercase		)		)
         __snake_case:						List[Any]			  =						reversed_block_out_channels[0]
         for i, up_block_type in enumerate(self.up_block_types		):
           __snake_case:						Any			  =						output_channel
           __snake_case:						Any			  =						reversed_block_out_channels[i]
           __snake_case:						Tuple			  =						reversed_block_out_channels[min(i + 1							,    len(__lowercase		) - 1		)]

           __snake_case:						Any			  =						i == len(__lowercase		) - 1

           if up_block_type == "CrossAttnUpBlock2D":
             __snake_case:						Tuple			  =						FlaxCrossAttnUpBlockaD(
                 in_channels=__lowercase							,    out_channels=__lowercase							,    prev_output_channel=__lowercase							,    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:						int			  =						FlaxUpBlockaD(
                 in_channels=__lowercase							,    out_channels=__lowercase							,    prev_output_channel=__lowercase							,    num_layers=self.layers_per_block + 1							,    add_upsample=not is_final_block							,    dropout=self.dropout							,    dtype=self.dtype							,    )

           up_blocks.append(__lowercase		)
           __snake_case:						Optional[int]			  =						output_channel
         __snake_case:						Optional[int]			  =						up_blocks

         # out
         __snake_case:						str			  =						nn.GroupNorm(num_groups=32							,    epsilon=1E-5		)
         __snake_case:						Tuple			  =						nn.Conv(
             self.out_channels							,    kernel_size=(3, 3)							,    strides=(1, 1)							,    padding=((1, 1), (1, 1))							,    dtype=self.dtype							,    )






       def __call__(  self			:   str							,    A			:   Tuple							,    A			:   Dict							,    A			:   Union[str, Any]							,    A			:   str=None							,    A			:   int=None							,    A			:   Optional[int] = True							,    A			:   Optional[Any] = False							,    ):
         # 1. time
         if not isinstance(__lowercase							,    jnp.ndarray		):
           __snake_case:						List[Any]			  =						jnp.array([timesteps]							,    dtype=jnp.intaa		)
         elif isinstance(__lowercase							,    jnp.ndarray		) and len(timesteps.shape		) == 0:
           __snake_case:						List[str]			  =						timesteps.astype(dtype=jnp.floataa		)
           __snake_case:						Any			  =						jnp.expand_dims(__lowercase							,    0		)

         __snake_case:						Dict			  =						self.time_proj(__lowercase		)
         __snake_case:						List[str]			  =						self.time_embedding(__lowercase		)

         # 2. pre-process
         __snake_case:						Optional[Any]			  =						jnp.transpose(__lowercase							,    (0, 2, 3, 1)		)
         __snake_case:						Optional[int]			  =						self.conv_in(__lowercase		)

         # 3. down
         __snake_case:						int			  =						(sample,)
         for down_block in self.down_blocks:
           if isinstance(__lowercase							,    __lowercase		):
             __snake_case:						str			  =						down_block(__lowercase							,    __lowercase							,    __lowercase							,    deterministic=not train		)
           else:
             __snake_case:						List[Any]			  =						down_block(__lowercase							,    __lowercase							,    deterministic=not train		)
           down_block_res_samples += res_samples

         if down_block_additional_residuals is not None:
           __snake_case:						Optional[int]			  =						()

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

           __snake_case:						int			  =						new_down_block_res_samples

         # 4. mid
         __snake_case:						List[str]			  =						self.mid_block(__lowercase							,    __lowercase							,    __lowercase							,    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:						Optional[int]			  =						down_block_res_samples[-(self.layers_per_block + 1) :]
           __snake_case:						Union[str, Any]			  =						down_block_res_samples[: -(self.layers_per_block + 1)]
           if isinstance(__lowercase							,    __lowercase		):
             __snake_case:						List[Any]			  =						up_block(
                 __lowercase							,    temb=__lowercase							,    encoder_hidden_states=__lowercase							,    res_hidden_states_tuple=__lowercase							,    deterministic=not train							,    )
           else:
             __snake_case:						List[str]			  =						up_block(__lowercase							,    temb=__lowercase							,    res_hidden_states_tuple=__lowercase							,    deterministic=not train		)

        # 6. post-process
         __snake_case:						int			  =						self.conv_norm_out(__lowercase		)
         __snake_case:						Optional[Any]			  =						nn.silu(__lowercase		)
         __snake_case:						Tuple			  =						self.conv_out(__lowercase		)
         __snake_case:						Optional[int]			  =						jnp.transpose(__lowercase							,    (0, 3, 1, 2)		)

         if not return_dict:
           return (sample,)

         return FlaxUNetaDConditionOutput(sample=__lowercase		)













import inspect
import unittest

from transformers import MobileViTConfig
from transformers.testing_utils import require_torch, 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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import MobileViTForImageClassification, MobileViTForSemanticSegmentation, MobileViTModel
      from transformers.models.mobilevit.modeling_mobilevit import MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import MobileViTImageProcessor







class    __snake_case	(							__lowerCamelCase      ):





       '''simple docstring'''







       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         __snake_case:						Optional[int]			  =						self.config_class(**self.inputs_dict		)
         self.parent.assertTrue(hasattr(A							,    """hidden_sizes"""		)		)
         self.parent.assertTrue(hasattr(A							,    """neck_hidden_sizes"""		)		)
         self.parent.assertTrue(hasattr(A							,    """num_attention_heads"""		)		)







class    __snake_case	:





       '''simple docstring'''







       def __init__(  self			:   int							,    A			:   str							,    A			:   Dict=13							,    A			:   str=32							,    A			:   Any=2							,    A			:   Optional[Any]=3							,    A			:   str=640							,    A			:   Tuple=4							,    A			:   Dict="silu"							,    A			:   List[Any]=3							,    A			:   Any=32							,    A			:   Any=0.1							,    A			:   int=0.1							,    A			:   Dict=0.1							,    A			:   Optional[Any]=0.02							,    A			:   List[Any]=True							,    A			:   Tuple=True							,    A			:   Any=10							,    A			:   Optional[int]=None							,    ):
         __snake_case:						List[Any]			  =						parent
         __snake_case:						Dict			  =						batch_size
         __snake_case:						int			  =						image_size
         __snake_case:						Tuple			  =						patch_size
         __snake_case:						Tuple			  =						num_channels
         __snake_case:						str			  =						last_hidden_size
         __snake_case:						Dict			  =						num_attention_heads
         __snake_case:						Dict			  =						hidden_act
         __snake_case:						Tuple			  =						conv_kernel_size
         __snake_case:						List[str]			  =						output_stride
         __snake_case:						List[str]			  =						hidden_dropout_prob
         __snake_case:						Optional[Any]			  =						attention_probs_dropout_prob
         __snake_case:						int			  =						classifier_dropout_prob
         __snake_case:						List[Any]			  =						use_labels
         __snake_case:						Union[str, Any]			  =						is_training
         __snake_case:						Union[str, Any]			  =						num_labels
         __snake_case:						str			  =						initializer_range
         __snake_case:						List[Any]			  =						scope




       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         __snake_case:						Union[str, Any]			  =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

         __snake_case:						Tuple			  =						None
         __snake_case:						Any			  =						None
         if self.use_labels:
           __snake_case:						Union[str, Any]			  =						ids_tensor([self.batch_size]							,    self.num_labels		)
           __snake_case:						str			  =						ids_tensor([self.batch_size, self.image_size, self.image_size]							,    self.num_labels		)

         __snake_case:						Any			  =						self.get_config()

         return config, pixel_values, labels, pixel_labels




       def 							UpperCAmelCase__			(  self			:   int		):
         return MobileViTConfig(
             image_size=self.image_size							,    patch_size=self.patch_size							,    num_channels=self.num_channels							,    num_attention_heads=self.num_attention_heads							,    hidden_act=self.hidden_act							,    conv_kernel_size=self.conv_kernel_size							,    output_stride=self.output_stride							,    hidden_dropout_prob=self.hidden_dropout_prob							,    attention_probs_dropout_prob=self.attention_probs_dropout_prob							,    classifier_dropout_prob=self.classifier_dropout_prob							,    initializer_range=self.initializer_range							,    )




       def 							UpperCAmelCase__			(  self			:   str							,    A			:   Optional[Any]							,    A			:   Any							,    A			:   Any							,    A			:   Union[str, Any]		):
         __snake_case:						List[Any]			  =						MobileViTModel(config=A		)
         model.to(A		)
         model.eval()
         __snake_case:						int			  =						model(A		)
         self.parent.assertEqual(
             result.last_hidden_state.shape							,    (
                 self.batch_size,
                 self.last_hidden_size,
                 self.image_size // self.output_stride,
                 self.image_size // self.output_stride,
             )							,    )




       def 							UpperCAmelCase__			(  self			:   str							,    A			:   List[Any]							,    A			:   Any							,    A			:   Any							,    A			:   int		):
         __snake_case:						str			  =						self.num_labels
         __snake_case:						Optional[int]			  =						MobileViTForImageClassification(A		)
         model.to(A		)
         model.eval()
         __snake_case:						Union[str, Any]			  =						model(A							,    labels=A		)
         self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.num_labels)		)




       def 							UpperCAmelCase__			(  self			:   Optional[int]							,    A			:   str							,    A			:   Optional[Any]							,    A			:   int							,    A			:   str		):
         __snake_case:						List[Any]			  =						self.num_labels
         __snake_case:						Dict			  =						MobileViTForSemanticSegmentation(A		)
         model.to(A		)
         model.eval()
         __snake_case:						Union[str, Any]			  =						model(A		)
         self.parent.assertEqual(
             result.logits.shape							,    (
                 self.batch_size,
                 self.num_labels,
                 self.image_size // self.output_stride,
                 self.image_size // self.output_stride,
             )							,    )
         __snake_case:						Tuple			  =						model(A							,    labels=A		)
         self.parent.assertEqual(
             result.logits.shape							,    (
                 self.batch_size,
                 self.num_labels,
                 self.image_size // self.output_stride,
                 self.image_size // self.output_stride,
             )							,    )






       def 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						Tuple			  =						self.prepare_config_and_inputs()
         __snake_case  ,       __snake_case  ,       __snake_case  ,       __snake_case:						Any			  =						config_and_inputs
         __snake_case:						Optional[int]			  =						{"""pixel_values""": pixel_values}
         return config, inputs_dict







@require_torch
class    __snake_case	(							__lowerCamelCase       ,      __lowerCamelCase       ,      unittest.TestCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					(
           (MobileViTModel, MobileViTForImageClassification, MobileViTForSemanticSegmentation)
           if is_torch_available()
           else ()
       )
       lowerCAmelCase__						=					(
           {
               """feature-extraction""": MobileViTModel,
               """image-classification""": MobileViTForImageClassification,
               """image-segmentation""": MobileViTForSemanticSegmentation,
           }
           if is_torch_available()
           else {}
       )

       lowerCAmelCase__						=					False
       lowerCAmelCase__						=					False
       lowerCAmelCase__						=					False
       lowerCAmelCase__						=					False




       def 							UpperCAmelCase__			(  self			:   List[str]		):
         __snake_case:						List[Any]			  =						MobileViTModelTester(self		)
         __snake_case:						str			  =						MobileViTConfigTester(self							,    config_class=A							,    has_text_modality=A		)




       def 							UpperCAmelCase__			(  self			:   str		):
         self.config_tester.run_common_tests()




       @unittest.skip(reason="""MobileViT does not use inputs_embeds"""		)
       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         pass




       @unittest.skip(reason="""MobileViT does not support input and output embeddings"""		)
       def 							UpperCAmelCase__			(  self			:   Dict		):
         pass




       @unittest.skip(reason="""MobileViT does not output attentions"""		)
       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         pass




       def 							UpperCAmelCase__			(  self			:   str		):
         __snake_case  ,       __snake_case:						Dict			  =						self.model_tester.prepare_config_and_inputs_for_common()

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

           __snake_case:						List[Any]			  =						["""pixel_values"""]
           self.assertListEqual(arg_names[:1]							,    A		)




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




       def 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						Optional[Any]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*A		)




       def 							UpperCAmelCase__			(  self			:   Dict		):
         def check_hidden_states_output(A			:   List[Any]							,    A			:   int							,    A			:   Tuple		):
           __snake_case:						List[str]			  =						model_class(A		)
           model.to(A		)
           model.eval()

           with torch.no_grad():
             __snake_case:						str			  =						model(**self._prepare_for_class(A							,    A		)		)

           __snake_case:						Optional[int]			  =						outputs.hidden_states

           __snake_case:						Any			  =						5
           self.assertEqual(len(A		)							,    A		)

           # MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
           # with the width and height being successively divided by 2.
           __snake_case:						Union[str, Any]			  =						2
           for i in range(len(A		)		):
             self.assertListEqual(
                 list(hidden_states[i].shape[-2:]		)							,    [self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]							,    )
             divisor *= 2

           self.assertEqual(self.model_tester.output_stride							,    divisor // 2		)

         __snake_case  ,       __snake_case:						List[Any]			  =						self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           __snake_case:						Optional[Any]			  =						True
           check_hidden_states_output(A							,    A							,    A		)

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           __snake_case:						Dict			  =						True

           check_hidden_states_output(A							,    A							,    A		)




       def 							UpperCAmelCase__			(  self			:   int		):
         __snake_case:						Optional[Any]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*A		)




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						Optional[int]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_semantic_segmentation(*A		)






       @slow
       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         for model_name in MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           __snake_case:						List[Any]			  =						MobileViTModel.from_pretrained(A		)
           self.assertIsNotNone(A		)







def       A__							(   )       ->   Optional[int]:
  __snake_case:						Optional[Any]			  =						Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""")
  return image







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





       '''simple docstring'''







       @cached_property
       def 							UpperCAmelCase__			(  self			:   Dict		):
         return MobileViTImageProcessor.from_pretrained("""apple/mobilevit-xx-small"""		) if is_vision_available() else None




       @slow
       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         __snake_case:						Tuple			  =						MobileViTForImageClassification.from_pretrained("""apple/mobilevit-xx-small"""		).to(A		)

         __snake_case:						str			  =						self.default_image_processor
         __snake_case:						Optional[Any]			  =						prepare_img()
         __snake_case:						List[Any]			  =						image_processor(images=A							,    return_tensors="""pt"""		).to(A		)

         # forward pass
         with torch.no_grad():
           __snake_case:						Dict			  =						model(**A		)

         # verify the logits
         __snake_case:						List[str]			  =						torch.Size((1, 1_000)		)
         self.assertEqual(outputs.logits.shape							,    A		)

         __snake_case:						Union[str, Any]			  =						torch.tensor([-1.9364, -1.2327, -0.4653]		).to(A		)

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




       @slow
       def 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						Tuple			  =						MobileViTForSemanticSegmentation.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""		)
         __snake_case:						List[str]			  =						model.to(A		)

         __snake_case:						Dict			  =						MobileViTImageProcessor.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""		)

         __snake_case:						List[Any]			  =						prepare_img()
         __snake_case:						List[str]			  =						image_processor(images=A							,    return_tensors="""pt"""		).to(A		)

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

         # verify the logits
         __snake_case:						Dict			  =						torch.Size((1, 21, 32, 32)		)
         self.assertEqual(logits.shape							,    A		)

         __snake_case:						Optional[int]			  =						torch.tensor(
             [
                 [[6.9713, 6.9786, 7.2422], [7.2893, 7.2825, 7.4446], [7.6580, 7.8797, 7.9420]],
                 [[-10.6869, -10.3250, -10.3471], [-10.4228, -9.9868, -9.7132], [-11.0405, -11.0221, -10.7318]],
                 [[-3.3089, -2.8539, -2.6740], [-3.2706, -2.5621, -2.5108], [-3.2534, -2.6615, -2.6651]],
             ]							,    device=A							,    )

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






       @slow
       def 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						int			  =						MobileViTForSemanticSegmentation.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""		)
         __snake_case:						str			  =						model.to(A		)

         __snake_case:						Optional[Any]			  =						MobileViTImageProcessor.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""		)

         __snake_case:						List[str]			  =						prepare_img()
         __snake_case:						Optional[int]			  =						image_processor(images=A							,    return_tensors="""pt"""		).to(A		)

         # forward pass
         with torch.no_grad():
           __snake_case:						Dict			  =						model(**A		)

         __snake_case:						List[Any]			  =						outputs.logits.detach().cpu()

         __snake_case:						List[str]			  =						image_processor.post_process_semantic_segmentation(outputs=A							,    target_sizes=[(50, 60)]		)
         __snake_case:						str			  =						torch.Size((50, 60)		)
         self.assertEqual(segmentation[0].shape							,    A		)

         __snake_case:						int			  =						image_processor.post_process_semantic_segmentation(outputs=A		)
         __snake_case:						Tuple			  =						torch.Size((32, 32)		)
         self.assertEqual(segmentation[0].shape							,    A		)






