


'''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 A__	(						UpperCAmelCase__				):






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




             '''simple docstring'''
             _a   :			List[Any]					 =self.config_class(**self.inputs_dict     )
             self.parent.assertTrue(hasattr(SCREAMING_SNAKE_CASE					,						"""hidden_sizes"""     )     )
             self.parent.assertTrue(hasattr(SCREAMING_SNAKE_CASE					,						"""neck_hidden_sizes"""     )     )
             self.parent.assertTrue(hasattr(SCREAMING_SNAKE_CASE					,						"""num_attention_heads"""     )     )







class A__	:






      def __init__(			self      :List[str]					,						SCREAMING_SNAKE_CASE      :Dict					,						SCREAMING_SNAKE_CASE      :Union[str, Any]=1_3					,						SCREAMING_SNAKE_CASE      :int=3_2					,						SCREAMING_SNAKE_CASE      :List[Any]=2					,						SCREAMING_SNAKE_CASE      :Optional[Any]=3					,						SCREAMING_SNAKE_CASE      :List[str]=6_4_0					,						SCREAMING_SNAKE_CASE      :Optional[Any]=4					,						SCREAMING_SNAKE_CASE      :Tuple="silu"					,						SCREAMING_SNAKE_CASE      :List[Any]=3					,						SCREAMING_SNAKE_CASE      :Dict=3_2					,						SCREAMING_SNAKE_CASE      :Union[str, Any]=0.1					,						SCREAMING_SNAKE_CASE      :int=0.1					,						SCREAMING_SNAKE_CASE      :List[str]=0.1					,						SCREAMING_SNAKE_CASE      :List[Any]=0.02					,						SCREAMING_SNAKE_CASE      :Any=True					,						SCREAMING_SNAKE_CASE      :List[str]=True					,						SCREAMING_SNAKE_CASE      :List[str]=1_0					,						SCREAMING_SNAKE_CASE      :str=None					,						)	-> Union[str, Any]:




             '''simple docstring'''
             _a   :			List[Any]					 =parent
             _a   :			Optional[Any]					 =batch_size
             _a   :			List[str]					 =image_size
             _a   :			Tuple					 =patch_size
             _a   :			str					 =num_channels
             _a   :			str					 =last_hidden_size
             _a   :			Optional[Any]					 =num_attention_heads
             _a   :			Optional[int]					 =hidden_act
             _a   :			Union[str, Any]					 =conv_kernel_size
             _a   :			Any					 =output_stride
             _a   :			Dict					 =hidden_dropout_prob
             _a   :			Optional[int]					 =attention_probs_dropout_prob
             _a   :			Tuple					 =classifier_dropout_prob
             _a   :			Optional[int]					 =use_labels
             _a   :			List[str]					 =is_training
             _a   :			Any					 =num_labels
             _a   :			Optional[int]					 =initializer_range
             _a   :			List[str]					 =scope






      def __UpperCAmelCase	(			self      :int     )	-> Any:




             '''simple docstring'''
             _a   :			Dict					 =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

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

             _a   :			List[str]					 =self.get_config()

             return config, pixel_values, labels, pixel_labels






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




             '''simple docstring'''
             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      :List[str]					,						SCREAMING_SNAKE_CASE      :Union[str, Any]					,						SCREAMING_SNAKE_CASE      :Dict					,						SCREAMING_SNAKE_CASE      :List[Any]					,						SCREAMING_SNAKE_CASE      :Dict     )	-> Dict:




             '''simple docstring'''
             _a   :			Any					 =TFMobileViTModel(config=SCREAMING_SNAKE_CASE     )
             _a   :			Tuple					 =model(SCREAMING_SNAKE_CASE					,						training=SCREAMING_SNAKE_CASE     )
             _a   :			Union[str, 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 __UpperCAmelCase	(			self      :Union[str, Any]					,						SCREAMING_SNAKE_CASE      :Dict					,						SCREAMING_SNAKE_CASE      :int					,						SCREAMING_SNAKE_CASE      :Tuple					,						SCREAMING_SNAKE_CASE      :str     )	-> List[Any]:




             '''simple docstring'''
             _a   :			Optional[int]					 =self.num_labels
             _a   :			List[Any]					 =TFMobileViTForImageClassification(SCREAMING_SNAKE_CASE     )
             _a   :			Optional[Any]					 =model(SCREAMING_SNAKE_CASE					,						labels=SCREAMING_SNAKE_CASE					,						training=SCREAMING_SNAKE_CASE     )
             self.parent.assertEqual(result.logits.shape					,						(self.batch_size, self.num_labels)     )






      def __UpperCAmelCase	(			self      :Dict					,						SCREAMING_SNAKE_CASE      :List[str]					,						SCREAMING_SNAKE_CASE      :List[str]					,						SCREAMING_SNAKE_CASE      :List[Any]					,						SCREAMING_SNAKE_CASE      :str     )	-> Dict:




             '''simple docstring'''
             _a   :			int					 =self.num_labels
             _a   :			Optional[int]					 =TFMobileViTForSemanticSegmentation(SCREAMING_SNAKE_CASE     )
             _a   :			Optional[int]					 =self.image_size // self.output_stride

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

             _a   :			Optional[int]					 =model(SCREAMING_SNAKE_CASE					,						labels=SCREAMING_SNAKE_CASE					,						training=SCREAMING_SNAKE_CASE     )
             self.parent.assertEqual(
                 result.logits.shape					,						(self.batch_size, self.num_labels, expected_height, expected_width)     )






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




             '''simple docstring'''
             _a   :			Dict					 =self.prepare_config_and_inputs()
             _a   :			Dict					 =config_and_inputs
             _a   :			Union[str, Any]					 ={"""pixel_values""": pixel_values}
             return config, inputs_dict







@require_tf
class A__	(						UpperCAmelCase__			,					UpperCAmelCase__			,					unittest.TestCase				):
      __UpperCamelCase	:  Dict     	=							(
          (TFMobileViTModel, TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation)
          if is_tf_available()
          else ()
      )
      __UpperCamelCase	:  List[str]     	=							(
          {"feature-extraction": TFMobileViTModel, "image-classification": TFMobileViTForImageClassification}
          if is_tf_available()
          else {}
      )

      __UpperCamelCase	:  Union[str, Any]     	=							False
      __UpperCamelCase	:  Optional[Any]     	=							False
      __UpperCamelCase	:  Optional[int]     	=							False
      __UpperCamelCase	:  List[str]     	=							False
      __UpperCamelCase	:  int     	=							False






      def __UpperCAmelCase	(			self      :Optional[int]     )	-> Union[str, Any]:




             '''simple docstring'''
             _a   :			List[str]					 =TFMobileViTModelTester(self     )
             _a   :			List[Any]					 =TFMobileViTConfigTester(self					,						config_class=SCREAMING_SNAKE_CASE					,						has_text_modality=SCREAMING_SNAKE_CASE     )






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




             '''simple docstring'''
             self.config_tester.run_common_tests()






      @unittest.skip(reason="""MobileViT does not use inputs_embeds"""     )
      def __UpperCAmelCase	(			self      :Any     )	-> Optional[Any]:




             '''simple docstring'''
             pass






      @unittest.skip(reason="""MobileViT does not support input and output embeddings"""     )
      def __UpperCAmelCase	(			self      :str     )	-> Tuple:




             '''simple docstring'''
             pass






      @unittest.skip(reason="""MobileViT does not output attentions"""     )
      def __UpperCAmelCase	(			self      :List[str]     )	-> Dict:




             '''simple docstring'''
             pass






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




             '''simple docstring'''
             _a   :			int					 =self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                    _a   :			List[str]					 =model_class(SCREAMING_SNAKE_CASE     )
                    _a   :			str					 =inspect.signature(model.call     )
                    # signature.parameters is an OrderedDict => so arg_names order is deterministic
                    _a   :			Dict					 =[*signature.parameters.keys()]

                    _a   :			List[str]					 =["""pixel_values"""]
                    self.assertListEqual(arg_names[:1]					,						SCREAMING_SNAKE_CASE     )






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




             '''simple docstring'''
             _a   :			Union[str, Any]					 =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE     )






      def __UpperCAmelCase	(			self      :Dict     )	-> List[str]:




             '''simple docstring'''
             def check_hidden_states_output(SCREAMING_SNAKE_CASE      :Optional[int]					,						SCREAMING_SNAKE_CASE      :str					,						SCREAMING_SNAKE_CASE      :Optional[Any]     ):
                    _a   :			Dict					 =model_class(SCREAMING_SNAKE_CASE     )

                    _a   :			List[str]					 =model(**self._prepare_for_class(SCREAMING_SNAKE_CASE					,						SCREAMING_SNAKE_CASE     )     )

                    _a   :			Any					 =outputs.hidden_states

                    _a   :			List[Any]					 =5
                    self.assertEqual(len(SCREAMING_SNAKE_CASE     )					,						SCREAMING_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.
                    _a   :			Any					 =2
                    for i in range(len(SCREAMING_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     )

             _a   :			List[str]					 =self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                    _a   :			List[Any]					 =True
                    check_hidden_states_output(SCREAMING_SNAKE_CASE					,						SCREAMING_SNAKE_CASE					,						SCREAMING_SNAKE_CASE     )

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

                    check_hidden_states_output(SCREAMING_SNAKE_CASE					,						SCREAMING_SNAKE_CASE					,						SCREAMING_SNAKE_CASE     )






      def __UpperCAmelCase	(			self      :Dict     )	-> Optional[int]:




             '''simple docstring'''
             _a   :			List[Any]					 =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE     )






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




             '''simple docstring'''
             _a   :			Optional[int]					 =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_semantic_segmentation(*SCREAMING_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 __UpperCAmelCase	(			self      :Union[str, Any]     )	-> Tuple:




             '''simple docstring'''
             super().test_dataset_conversion()






      def __UpperCAmelCase	(			self      :List[str]					,						SCREAMING_SNAKE_CASE      :List[str]					,						SCREAMING_SNAKE_CASE      :Union[str, Any]					,						SCREAMING_SNAKE_CASE      :str=2e-1					,						SCREAMING_SNAKE_CASE      :Tuple=2e-1     )	-> Union[str, Any]:




             '''simple docstring'''
             self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE					,						SCREAMING_SNAKE_CASE					,						atol=SCREAMING_SNAKE_CASE					,						rtol=SCREAMING_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 __UpperCAmelCase	(			self      :Optional[Any]     )	-> Tuple:




             '''simple docstring'''
             _a   :			Any					 =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":
                           _a   :			Tuple					 =model_class(SCREAMING_SNAKE_CASE     )
                           if getattr(SCREAMING_SNAKE_CASE					,						"""hf_compute_loss"""					,						SCREAMING_SNAKE_CASE     ):
                                  super().test_keras_fit()






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




             '''simple docstring'''
             _a   :			List[Any]					 =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":
                           _a   :			Dict					 =5

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

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

                           _a   :			Any					 =model(SCREAMING_SNAKE_CASE					,						**SCREAMING_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
                           _a   :			Tuple					 =self._prepare_for_class(inputs_dict.copy()					,						SCREAMING_SNAKE_CASE					,						return_labels=SCREAMING_SNAKE_CASE     )
                           _a   :			Optional[int]					 ={"""input_ids""", """pixel_values""", """input_features"""}
                           _a   :			List[Any]					 =possible_input_names.intersection(set(SCREAMING_SNAKE_CASE     )     ).pop()
                           _a   :			Dict					 =prepared_for_class.pop(SCREAMING_SNAKE_CASE     )
                           if "labels" in prepared_for_class:
                                  _a   :			Tuple					 =prepared_for_class["""labels"""].numpy()
                                  if len(labels.shape     ) > 1 and labels.shape[1] != 1:
                                         # labels[0] = -100
                                         _a   :			Tuple					 =tf.convert_to_tensor(SCREAMING_SNAKE_CASE     )
                                         _a   :			Optional[Any]					 =model(SCREAMING_SNAKE_CASE					,						**SCREAMING_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
                           _a   :			Optional[Any]					 =self._prepare_for_class(inputs_dict.copy()					,						SCREAMING_SNAKE_CASE					,						return_labels=SCREAMING_SNAKE_CASE     )
                           _a   :			List[str]					 =model(SCREAMING_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
                           _a   :			Optional[Any]					 =self._prepare_for_class(inputs_dict.copy()					,						SCREAMING_SNAKE_CASE					,						return_labels=SCREAMING_SNAKE_CASE     )

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

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

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

                           for index, value in sorted_tuple_index_mapping:
                                  _a   :			Union[str, Any]					 =prepared_for_class[value]

                           _a   :			List[Any]					 =tuple(SCREAMING_SNAKE_CASE     )

                           # Send to model
                           _a   :			List[Any]					 =model(tuple_input[:-1]     )[0]

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






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




             '''simple docstring'''
             for model_name in TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _a   :			int					 =TFMobileViTModel.from_pretrained(SCREAMING_SNAKE_CASE     )
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE     )


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







@require_tf
class A__	(						unittest.TestCase				):






      @slow
      def __UpperCAmelCase	(			self      :Tuple     )	-> List[str]:




             '''simple docstring'''
             _a   :			Dict					 =TFMobileViTForImageClassification.from_pretrained("""apple/mobilevit-xx-small"""     )

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

             # forward pass
             _a   :			int					 =model(**SCREAMING_SNAKE_CASE					,						training=SCREAMING_SNAKE_CASE     )

             # verify the logits
             _a   :			int					 =tf.TensorShape((1, 1_0_0_0)     )
             self.assertEqual(outputs.logits.shape					,						SCREAMING_SNAKE_CASE     )

             _a   :			int					 =tf.constant([-1.9_364, -1.2_327, -0.4_653]     )

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






      @slow
      def __UpperCAmelCase	(			self      :str     )	-> str:




             '''simple docstring'''
             _a   :			Any					 =TFMobileViTForSemanticSegmentation.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""     )

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

             _a   :			int					 =prepare_img()
             _a   :			Any					 =image_processor(images=SCREAMING_SNAKE_CASE					,						return_tensors="""tf"""     )

             # forward pass
             _a   :			Tuple					 =model(inputs.pixel_values					,						training=SCREAMING_SNAKE_CASE     )
             _a   :			Optional[int]					 =outputs.logits

             # verify the logits
             _a   :			Optional[int]					 =tf.TensorShape((1, 2_1, 3_2, 3_2)     )
             self.assertEqual(logits.shape					,						SCREAMING_SNAKE_CASE     )

             _a   :			str					 =tf.constant(
                 [
                     [[6.9_713, 6.9_786, 7.2_422], [7.2_893, 7.2_825, 7.4_446], [7.6_580, 7.8_797, 7.9_420]],
                     [[-10.6_869, -10.3_250, -10.3_471], [-10.4_228, -9.9_868, -9.7_132], [-11.0_405, -11.0_221, -10.7_318]],
                     [[-3.3_089, -2.8_539, -2.6_740], [-3.2_706, -2.5_621, -2.5_108], [-3.2_534, -2.6_615, -2.6_651]],
                 ]     )

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






'''simple docstring'''







from typing import Dict, Optional, Tuple

import torch


def      SCREAMING_SNAKE_CASE_ (  _UpperCAmelCase     :     torch.Tensor  )  ->      torch.Tensor:
       _a   :			Tuple					 =boundaries[1] - boundaries[0]
       _a   :			Dict					 =boundaries + step / 2
       _a   :			List[Any]					 =torch.cat([bin_centers, (bin_centers[-1] + step).unsqueeze(-1  )]   ,dim=0  )
       return bin_centers


def      SCREAMING_SNAKE_CASE_ (  _UpperCAmelCase     :     torch.Tensor   ,_UpperCAmelCase     :     torch.Tensor   ,)  ->      Tuple[torch.Tensor, torch.Tensor]:
       _a   :			Any					 =_calculate_bin_centers(_UpperCAmelCase  )
       return (
           torch.sum(aligned_distance_error_probs * bin_centers   ,dim=-1  ),
           bin_centers[-1],
       )


def      SCREAMING_SNAKE_CASE_ (  _UpperCAmelCase     :     torch.Tensor   ,_UpperCAmelCase     :     int = 31   ,_UpperCAmelCase     :     int = 64   ,**_UpperCAmelCase     :     Dict   ,)  ->      Dict[str, torch.Tensor]:
       _a   :			Dict					 =torch.linspace(0   ,_UpperCAmelCase   ,steps=(no_bins - 1)   ,device=logits.device  )

       _a   :			int					 =torch.nn.functional.softmax(_UpperCAmelCase   ,dim=-1  )
       _a ,   _a   :			List[Any]					 =_calculate_expected_aligned_error(
           alignment_confidence_breaks=_UpperCAmelCase   ,aligned_distance_error_probs=_UpperCAmelCase   ,)

       return {
           "aligned_confidence_probs": aligned_confidence_probs,
           "predicted_aligned_error": predicted_aligned_error,
           "max_predicted_aligned_error": max_predicted_aligned_error,
       }


def      SCREAMING_SNAKE_CASE_ (  _UpperCAmelCase     :     torch.Tensor   ,_UpperCAmelCase     :     Optional[torch.Tensor] = None   ,_UpperCAmelCase     :     int = 31   ,_UpperCAmelCase     :     int = 64   ,_UpperCAmelCase     :     float = 1e-8   ,**_UpperCAmelCase     :     Optional[Any]   ,)  ->      torch.Tensor:
       if residue_weights is None:
              _a   :			Dict					 =logits.new_ones(logits.shape[-2]  )

       _a   :			str					 =torch.linspace(0   ,_UpperCAmelCase   ,steps=(no_bins - 1)   ,device=logits.device  )

       _a   :			Optional[int]					 =_calculate_bin_centers(_UpperCAmelCase  )
       torch.sum(_UpperCAmelCase  )
       _a   :			List[str]					 =logits.shape[-2]
       _a   :			List[Any]					 =max(_UpperCAmelCase   ,19  )

       _a   :			Union[str, Any]					 =1.2_4 * (clipped_n - 15) ** (1.0 / 3) - 1.8

       _a   :			Tuple					 =torch.nn.functional.softmax(_UpperCAmelCase   ,dim=-1  )

       _a   :			Dict					 =1.0 / (1 + (bin_centers**2) / (da**2))
       _a   :			int					 =torch.sum(probs * tm_per_bin   ,dim=-1  )

       _a   :			List[Any]					 =residue_weights / (eps + residue_weights.sum())
       _a   :			List[str]					 =torch.sum(predicted_tm_term * normed_residue_mask   ,dim=-1  )

       _a   :			int					 =per_alignment * residue_weights

       _a   :			List[Any]					 =(weighted == torch.max(_UpperCAmelCase  )).nonzero()[0]
       return per_alignment[tuple(_UpperCAmelCase  )]


