

import math
from typing import Optional

import torch
import torch.utils.checkpoint as checkpoint
from torch import nn

from ....modeling_utils import PreTrainedModel
from ....utils import add_start_docstrings, logging
from ...bert.modeling_bert import BertModel
from .configuration_retribert import RetriBertConfig


A__							=				logging.get_logger(__name__)

A__							=				[
    """yjernite/retribert-base-uncased""",
    # See all RetriBert models at https://huggingface.co/models?filter=retribert
]

class 					__UpperCamelCase					(       _snake_case      ):
	_lowercase	: Any												=	RetriBertConfig
	_lowercase	: Union[str, Any]												=	None
	_lowercase	: str												=	"""retribert"""







	def 						_SCREAMING_SNAKE_CASE				( self:			List[str]	,			__UpperCamelCase:			Tuple   ):

						'''simple docstring'''





						if isinstance(snake_case_	,			nn.Linear   ):
											module.weight.data.normal_(mean=0.0	,			std=self.config.initializer_range   )
											if module.bias is not None:
																module.bias.data.zero_()
						elif isinstance(snake_case_	,			nn.Embedding   ):
											module.weight.data.normal_(mean=0.0	,			std=self.config.initializer_range   )
											if module.padding_idx is not None:
																module.weight.data[module.padding_idx].zero_()
						elif isinstance(snake_case_	,			nn.LayerNorm   ):
											module.bias.data.zero_()
											module.weight.data.fill_(1.0   )




A__							=				r"""

    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
    etc.)

    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
    and behavior.

    Parameters:
        config ([`RetriBertConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

@add_start_docstrings(
    "Bert Based model to embed queries or document for document retrieval."	,    _snake_case	,    )
class 					__UpperCamelCase					(       _snake_case      ):







	def __init__( self:			Any	,			__UpperCamelCase:			Union[str, Any]   ):

						'''simple docstring'''





						super().__init__(snake_case_   )
						__magic_name__				=     config.projection_dim

						__magic_name__				=     BertModel(snake_case_   )
						__magic_name__				=     None if config.share_encoders else BertModel(snake_case_   )
						__magic_name__				=     nn.Dropout(config.hidden_dropout_prob   )
						__magic_name__				=     nn.Linear(config.hidden_size	,			config.projection_dim	,			bias=snake_case_   )
						__magic_name__				=     nn.Linear(config.hidden_size	,			config.projection_dim	,			bias=snake_case_   )

						__magic_name__				=     nn.CrossEntropyLoss(reduction='mean'   )

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







	def 						_SCREAMING_SNAKE_CASE				( self:			int	,			__UpperCamelCase:			Any	,			__UpperCamelCase:			Union[str, Any]	,			__UpperCamelCase:			str	,			__UpperCamelCase:			Optional[int]=-1	,			):

						'''simple docstring'''





						if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
											return sent_encoder(snake_case_	,			attention_mask=snake_case_   )[1]
						else:
											# prepare implicit variables
											__magic_name__				=     input_ids.device
											__magic_name__				=     input_ids.size()
											__magic_name__				=     torch.zeros(snake_case_	,			dtype=torch.long	,			device=snake_case_   )
											__magic_name__				=     [None] * sent_encoder.config.num_hidden_layers
											__magic_name__				=     sent_encoder.get_extended_attention_mask(
											    snake_case_	,			snake_case_   )

											# define function for checkpointing
											def partial_encode(*__UpperCamelCase:			str   ):
																__magic_name__				=     sent_encoder.encoder(
																    inputs[0]	,			attention_mask=inputs[1]	,			head_mask=snake_case_	,			)
																__magic_name__				=     encoder_outputs[0]
																__magic_name__				=     sent_encoder.pooler(snake_case_   )
																return pooled_output

											# run embedding layer on everything at once
											__magic_name__				=     sent_encoder.embeddings(
											    input_ids=snake_case_	,			position_ids=snake_case_	,			token_type_ids=snake_case_	,			inputs_embeds=snake_case_   )
											# run encoding and pooling on one mini-batch at a time
											__magic_name__				=     []
											for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size   )   ):
																__magic_name__				=     embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
																__magic_name__				=     extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
																__magic_name__				=     checkpoint.checkpoint(snake_case_	,			snake_case_	,			snake_case_   )
																pooled_output_list.append(snake_case_   )
											return torch.cat(snake_case_	,			dim=0   )







	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[Any]	,			__UpperCamelCase:			Any	,			__UpperCamelCase:			Any=None	,			__UpperCamelCase:			List[Any]=-1	,			):

						'''simple docstring'''





						__magic_name__				=     self.embed_sentences_checkpointed(
						    snake_case_	,			snake_case_	,			self.bert_query	,			snake_case_	,			)
						return self.project_query(snake_case_   )







	def 						_SCREAMING_SNAKE_CASE				( self:			int	,			__UpperCamelCase:			Tuple	,			__UpperCamelCase:			List[Any]=None	,			__UpperCamelCase:			List[str]=-1	,			):

						'''simple docstring'''





						__magic_name__				=     self.embed_sentences_checkpointed(
						    snake_case_	,			snake_case_	,			self.bert_query if self.bert_doc is None else self.bert_doc	,			snake_case_	,			)
						return self.project_doc(snake_case_   )







	def 						_SCREAMING_SNAKE_CASE				( self:			List[str]	,			__UpperCamelCase:			List[Any]	,			__UpperCamelCase:			Tuple	,			__UpperCamelCase:			Tuple	,			__UpperCamelCase:			List[Any]	,			__UpperCamelCase:			int = -1	,			):

						'''simple docstring'''





						__magic_name__				=     input_ids_query.device
						__magic_name__				=     self.embed_questions(snake_case_	,			snake_case_	,			snake_case_   )
						__magic_name__				=     self.embed_answers(snake_case_	,			snake_case_	,			snake_case_   )
						__magic_name__				=     torch.mm(snake_case_	,			a_reps.t()   )
						__magic_name__				=     self.ce_loss(snake_case_	,			torch.arange(compare_scores.shape[1]   ).to(snake_case_   )   )
						__magic_name__				=     self.ce_loss(compare_scores.t()	,			torch.arange(compare_scores.shape[0]   ).to(snake_case_   )   )
						__magic_name__				=     (loss_qa + loss_aq) / 2
						return loss






"""simple docstring"""
from __future__ import annotations

import inspect
import unittest

from transformers import MobileViTConfig
from transformers.file_utils import is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import numpy as np
    import tensorflow as tf

    from transformers import TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation, TFMobileViTModel
    from transformers.models.mobilevit.modeling_tf_mobilevit import TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import MobileViTImageProcessor






class   _UpperCAmelCase					(       _snake_case):
    def        lowerCamelCase__						(     self  ):
      _snake_case  :    List[str]            = self.config_class(**self.inputs_dict  )
      self.parent.assertTrue(hasattr(snake_case_    ,     "hidden_sizes"  )  )
      self.parent.assertTrue(hasattr(snake_case_    ,     "neck_hidden_sizes"  )  )
      self.parent.assertTrue(hasattr(snake_case_    ,     "num_attention_heads"  )  )






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

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

      _snake_case  :    Optional[Any]            = self.get_config()

      return config, pixel_values, labels, pixel_labels
    def        lowerCamelCase__						(     self  ):
      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        lowerCamelCase__						(     self    ,     snake_case_    ,     snake_case_    ,     snake_case_    ,     snake_case_  ):
      _snake_case  :    str            = TFMobileViTModel(config=snake_case_  )
      _snake_case  :    Optional[int]            = model(snake_case_    ,     training=snake_case_  )
      _snake_case  :    List[Any]            = self.image_size // self.output_stride
      self.parent.assertEqual(
          result.last_hidden_state.shape    ,     (self.batch_size, self.last_hidden_size, expected_height, expected_width)  )
    def        lowerCamelCase__						(     self    ,     snake_case_    ,     snake_case_    ,     snake_case_    ,     snake_case_  ):
      _snake_case  :    int            = self.num_labels
      _snake_case  :    Optional[Any]            = TFMobileViTForImageClassification(snake_case_  )
      _snake_case  :    str            = model(snake_case_    ,     labels=snake_case_    ,     training=snake_case_  )
      self.parent.assertEqual(result.logits.shape    ,     (self.batch_size, self.num_labels)  )
    def        lowerCamelCase__						(     self    ,     snake_case_    ,     snake_case_    ,     snake_case_    ,     snake_case_  ):
      _snake_case  :    List[str]            = self.num_labels
      _snake_case  :    Any            = TFMobileViTForSemanticSegmentation(snake_case_  )
      _snake_case  :    Optional[Any]            = self.image_size // self.output_stride

      _snake_case  :    Any            = model(snake_case_    ,     training=snake_case_  )
      self.parent.assertEqual(
          result.logits.shape    ,     (self.batch_size, self.num_labels, expected_height, expected_width)  )

      _snake_case  :    List[Any]            = model(snake_case_    ,     labels=snake_case_    ,     training=snake_case_  )
      self.parent.assertEqual(
          result.logits.shape    ,     (self.batch_size, self.num_labels, expected_height, expected_width)  )


    def        lowerCamelCase__						(     self  ):
      _snake_case  :    Union[str, Any]            = self.prepare_config_and_inputs()
      _snake_case      ,       _snake_case      ,       _snake_case      ,       _snake_case  :    str            = config_and_inputs
      _snake_case  :    Optional[int]            = {"pixel_values": pixel_values}
      return config, inputs_dict






@require_tf
class   _UpperCAmelCase					(       _snake_case		,       _snake_case		,       unittest.TestCase):
    __lowercase				:			Tuple	=       (
        (TFMobileViTModel, TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation)
        if is_tf_available()
        else ()
    )
    __lowercase				:			int	=       (
        {"""feature-extraction""": TFMobileViTModel, """image-classification""": TFMobileViTForImageClassification}
        if is_tf_available()
        else {}
    )

    __lowercase				:			Optional[int]	=       False
    __lowercase				:			Optional[int]	=       False
    __lowercase				:			List[Any]	=       False
    __lowercase				:			int	=       False
    __lowercase				:			Union[str, Any]	=       False
    def        lowerCamelCase__						(     self  ):
      _snake_case  :    Tuple            = TFMobileViTModelTester(self  )
      _snake_case  :    Dict            = TFMobileViTConfigTester(self    ,     config_class=snake_case_    ,     has_text_modality=snake_case_  )
    def        lowerCamelCase__						(     self  ):
      self.config_tester.run_common_tests()
    @unittest.skip(reason="MobileViT does not use inputs_embeds"  )
    def        lowerCamelCase__						(     self  ):
      pass
    @unittest.skip(reason="MobileViT does not support input and output embeddings"  )
    def        lowerCamelCase__						(     self  ):
      pass
    @unittest.skip(reason="MobileViT does not output attentions"  )
    def        lowerCamelCase__						(     self  ):
      pass
    def        lowerCamelCase__						(     self  ):
      _snake_case      ,       _snake_case  :    Optional[int]            = self.model_tester.prepare_config_and_inputs_for_common()

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

        _snake_case  :    List[str]            = ["pixel_values"]
        self.assertListEqual(arg_names[:1]    ,     snake_case_  )
    def        lowerCamelCase__						(     self  ):
      _snake_case  :    List[Any]            = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*snake_case_  )
    def        lowerCamelCase__						(     self  ):
      def check_hidden_states_output(snake_case_    ,     snake_case_    ,     snake_case_  ):
        _snake_case  :    Optional[int]            = model_class(snake_case_  )

        _snake_case  :    Dict            = model(**self._prepare_for_class(snake_case_    ,     snake_case_  )  )

        _snake_case  :    Optional[Any]            = outputs.hidden_states

        _snake_case  :    Any            = 5
        self.assertEqual(len(snake_case_  )    ,     snake_case_  )

        # 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  :    Dict            = 2
        for i in range(len(snake_case_  )  ):
          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  :    Optional[int]            = self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        _snake_case  :    Union[str, Any]            = True
        check_hidden_states_output(snake_case_    ,     snake_case_    ,     snake_case_  )

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

        check_hidden_states_output(snake_case_    ,     snake_case_    ,     snake_case_  )
    def        lowerCamelCase__						(     self  ):
      _snake_case  :    List[str]            = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*snake_case_  )
    def        lowerCamelCase__						(     self  ):
      _snake_case  :    Tuple            = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_semantic_segmentation(*snake_case_  )
    @unittest.skipIf(
        not is_tf_available() or len(tf.config.list_physical_devices("GPU"  )  ) == 0    ,     reason="TF does not support backprop for grouped convolutions on CPU."    ,     )
    def        lowerCamelCase__						(     self  ):
      super().test_dataset_conversion()
    def        lowerCamelCase__						(     self    ,     snake_case_    ,     snake_case_    ,     snake_case_=2E-1    ,     snake_case_=2E-1  ):
      self.assertTrue(np.allclose(snake_case_    ,     snake_case_    ,     atol=snake_case_    ,     rtol=snake_case_  )  )
    @unittest.skipIf(
        not is_tf_available() or len(tf.config.list_physical_devices("GPU"  )  ) == 0    ,     reason="TF does not support backprop for grouped convolutions on CPU."    ,     )
    @slow
    def        lowerCamelCase__						(     self  ):
      _snake_case      ,       _snake_case  :    Optional[int]            = self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        # Since `TFMobileViTModel` cannot operate with the default `fit()` method.
        if model_class.__name__ != "TFMobileViTModel":
          _snake_case  :    List[str]            = model_class(snake_case_  )
          if getattr(snake_case_    ,     "hf_compute_loss"    ,     snake_case_  ):
            super().test_keras_fit()
    def        lowerCamelCase__						(     self  ):
      _snake_case      ,       _snake_case  :    Optional[int]            = self.model_tester.prepare_config_and_inputs_for_common()
      for model_class in self.all_model_classes:
        # set an ignore index to correctly test the masked loss used in
        # `TFMobileViTForSemanticSegmentation`.
        if model_class.__name__ != "TFMobileViTForSemanticSegmentation":
          _snake_case  :    Optional[Any]            = 5

        _snake_case  :    Optional[Any]            = model_class(snake_case_  )
        if getattr(snake_case_    ,     "hf_compute_loss"    ,     snake_case_  ):
          # The number of elements in the loss should be the same as the number of elements in the label
          _snake_case  :    Optional[Any]            = self._prepare_for_class(inputs_dict.copy()    ,     snake_case_    ,     return_labels=snake_case_  )
          _snake_case  :    int            = prepared_for_class[
              sorted(prepared_for_class.keys() - inputs_dict.keys()    ,     reverse=snake_case_  )[0]
          ]
          _snake_case  :    Optional[Any]            = added_label.shape.as_list()[:1]

          # Test that model correctly compute the loss with kwargs
          _snake_case  :    List[str]            = self._prepare_for_class(inputs_dict.copy()    ,     snake_case_    ,     return_labels=snake_case_  )
          _snake_case  :    str            = {"input_ids", "pixel_values", "input_features"}
          _snake_case  :    Optional[int]            = possible_input_names.intersection(set(snake_case_  )  ).pop()
          _snake_case  :    int            = prepared_for_class.pop(snake_case_  )

          _snake_case  :    Union[str, Any]            = model(snake_case_    ,     **snake_case_  )[0]
          self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]  )

          # Test that model correctly compute the loss when we mask some positions
          _snake_case  :    Optional[Any]            = self._prepare_for_class(inputs_dict.copy()    ,     snake_case_    ,     return_labels=snake_case_  )
          _snake_case  :    Any            = {"input_ids", "pixel_values", "input_features"}
          _snake_case  :    Optional[int]            = possible_input_names.intersection(set(snake_case_  )  ).pop()
          _snake_case  :    List[str]            = prepared_for_class.pop(snake_case_  )
          if "labels" in prepared_for_class:
            _snake_case  :    Optional[int]            = prepared_for_class["labels"].numpy()
            if len(labels.shape  ) > 1 and labels.shape[1] != 1:
              # labels[0] = -100
              _snake_case  :    Optional[int]            = tf.convert_to_tensor(snake_case_  )
              _snake_case  :    Union[str, Any]            = model(snake_case_    ,     **snake_case_  )[0]
              self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]  )
              self.assertTrue(not np.any(np.isnan(loss.numpy()  )  )  )

                # Test that model correctly compute the loss with a dict
          _snake_case  :    Tuple            = self._prepare_for_class(inputs_dict.copy()    ,     snake_case_    ,     return_labels=snake_case_  )
          _snake_case  :    Tuple            = model(snake_case_  )[0]
          self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]  )

          # Test that model correctly compute the loss with a tuple
          _snake_case  :    Dict            = self._prepare_for_class(inputs_dict.copy()    ,     snake_case_    ,     return_labels=snake_case_  )

          # Get keys that were added with the _prepare_for_class function
          _snake_case  :    Union[str, Any]            = prepared_for_class.keys() - inputs_dict.keys()
          _snake_case  :    Any            = inspect.signature(model.call  ).parameters
          _snake_case  :    List[Any]            = list(signature.keys()  )

          # Create a dictionary holding the location of the tensors in the tuple
          _snake_case  :    Union[str, Any]            = {0: input_name}
          for label_key in label_keys:
            _snake_case  :    Tuple            = signature_names.index(snake_case_  )
            _snake_case  :    int            = label_key
          _snake_case  :    List[Any]            = sorted(tuple_index_mapping.items()  )
          # Initialize a list with their default values, update the values and convert to a tuple
          _snake_case  :    List[Any]            = []

          for name in signature_names:
            if name != "kwargs":
              list_input.append(signature[name].default  )

          for index, value in sorted_tuple_index_mapping:
            _snake_case  :    int            = prepared_for_class[value]

          _snake_case  :    str            = tuple(snake_case_  )

          # Send to model
          _snake_case  :    Union[str, Any]            = model(tuple_input[:-1]  )[0]

          self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]  )


    @slow
    def        lowerCamelCase__						(     self  ):
      for model_name in TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _snake_case  :    Optional[int]            = TFMobileViTModel.from_pretrained(snake_case_  )
        self.assertIsNotNone(snake_case_  )


def 					a__ ( ):

  """simple docstring"""







  _snake_case  :    Union[str, Any]            = Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"		)
  return image






@require_tf
class   _UpperCAmelCase					(       unittest.TestCase):
    @slow
    def        lowerCamelCase__						(     self  ):
      _snake_case  :    Optional[Any]            = TFMobileViTForImageClassification.from_pretrained("apple/mobilevit-xx-small"  )

      _snake_case  :    List[Any]            = MobileViTImageProcessor.from_pretrained("apple/mobilevit-xx-small"  )
      _snake_case  :    List[Any]            = prepare_img()
      _snake_case  :    List[str]            = image_processor(images=snake_case_    ,     return_tensors="tf"  )

      # forward pass
      _snake_case  :    Union[str, Any]            = model(**snake_case_    ,     training=snake_case_  )

      # verify the logits
      _snake_case  :    Optional[Any]            = tf.TensorShape((1, 10_00)  )
      self.assertEqual(outputs.logits.shape    ,     snake_case_  )

      _snake_case  :    List[Any]            = tf.constant([-1.9364, -1.2327, -0.4653]  )

      tf.debugging.assert_near(outputs.logits[0, :3]    ,     snake_case_    ,     atol=1E-4    ,     rtol=1E-04  )


    @slow
    def        lowerCamelCase__						(     self  ):
      # `from_pt` will be removed
      _snake_case  :    List[str]            = TFMobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-xx-small"  )

      _snake_case  :    int            = MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small"  )

      _snake_case  :    List[Any]            = prepare_img()
      _snake_case  :    str            = image_processor(images=snake_case_    ,     return_tensors="tf"  )

      # forward pass
      _snake_case  :    Tuple            = model(inputs.pixel_values    ,     training=snake_case_  )
      _snake_case  :    Tuple            = outputs.logits

      # verify the logits
      _snake_case  :    Any            = tf.TensorShape((1, 21, 32, 32)  )
      self.assertEqual(logits.shape    ,     snake_case_  )

      _snake_case  :    str            = tf.constant(
          [
              [[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]],
          ]  )

      tf.debugging.assert_near(logits[0, :3, :3, :3]    ,     snake_case_    ,     rtol=1E-4    ,     atol=1E-4  )

