


import collections
import inspect
import unittest

from transformers import DinatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch
   from torch import nn

   from transformers import DinatBackbone, DinatForImageClassification, DinatModel
   from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
   from PIL import Image

   from transformers import AutoImageProcessor






class        UpperCAmelCase__			:



    """simple docstring"""






    def __init__(     self   ,  A_   ,  A_=13   ,  A_=64   ,  A_=4   ,  A_=3   ,  A_=16   ,  A_=[1, 2, 1]   ,  A_=[2, 4, 8]   ,  A_=3   ,  A_=[[3], [1, 2], [1]]   ,  A_=2.0   ,  A_=True   ,  A_=0.0   ,  A_=0.0   ,  A_=0.1   ,  A_="gelu"   ,  A_=True   ,  A_=0.02   ,  A_=1E-5   ,  A_=True   ,  A_=None   ,  A_=True   ,  A_=10   ,  A_=["stage1", "stage2"]   ,  A_=[1, 2]   ,  )    ->			Dict:
         __UpperCamelCase       				=parent
         __UpperCamelCase       				=batch_size
         __UpperCamelCase       				=image_size
         __UpperCamelCase       				=patch_size
         __UpperCamelCase       				=num_channels
         __UpperCamelCase       				=embed_dim
         __UpperCamelCase       				=depths
         __UpperCamelCase       				=num_heads
         __UpperCamelCase       				=kernel_size
         __UpperCamelCase       				=dilations
         __UpperCamelCase       				=mlp_ratio
         __UpperCamelCase       				=qkv_bias
         __UpperCamelCase       				=hidden_dropout_prob
         __UpperCamelCase       				=attention_probs_dropout_prob
         __UpperCamelCase       				=drop_path_rate
         __UpperCamelCase       				=hidden_act
         __UpperCamelCase       				=patch_norm
         __UpperCamelCase       				=layer_norm_eps
         __UpperCamelCase       				=initializer_range
         __UpperCamelCase       				=is_training
         __UpperCamelCase       				=scope
         __UpperCamelCase       				=use_labels
         __UpperCamelCase       				=num_labels
         __UpperCamelCase       				=out_features
         __UpperCamelCase       				=out_indices






    def _a   (     self       )    ->			List[Any]:
         __UpperCamelCase       				=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

         __UpperCamelCase       				=None
         if self.use_labels:
              __UpperCamelCase       				=ids_tensor([self.batch_size]   ,  self.num_labels       )

         __UpperCamelCase       				=self.get_config()

         return config, pixel_values, labels






    def _a   (     self       )    ->			List[str]:
         return DinatConfig(
             num_labels=self.num_labels   ,  image_size=self.image_size   ,  patch_size=self.patch_size   ,  num_channels=self.num_channels   ,  embed_dim=self.embed_dim   ,  depths=self.depths   ,  num_heads=self.num_heads   ,  kernel_size=self.kernel_size   ,  dilations=self.dilations   ,  mlp_ratio=self.mlp_ratio   ,  qkv_bias=self.qkv_bias   ,  hidden_dropout_prob=self.hidden_dropout_prob   ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,  drop_path_rate=self.drop_path_rate   ,  hidden_act=self.hidden_act   ,  patch_norm=self.patch_norm   ,  layer_norm_eps=self.layer_norm_eps   ,  initializer_range=self.initializer_range   ,  out_features=self.out_features   ,  out_indices=self.out_indices   ,  )






    def _a   (     self   ,  A_   ,  A_   ,  A_       )    ->			int:
         __UpperCamelCase       				=DinatModel(config=A_       )
         model.to(A_       )
         model.eval()
         __UpperCamelCase       				=model(A_       )

         __UpperCamelCase       				=__UpperCamelCase       				=(config.image_size // config.patch_size) // (2 ** (len(config.depths       ) - 1))
         __UpperCamelCase       				=int(config.embed_dim * 2 ** (len(config.depths       ) - 1)       )

         self.parent.assertEqual(
             result.last_hidden_state.shape   ,  (self.batch_size, expected_height, expected_width, expected_dim)       )






    def _a   (     self   ,  A_   ,  A_   ,  A_       )    ->			Any:
         __UpperCamelCase       				=DinatForImageClassification(A_       )
         model.to(A_       )
         model.eval()
         __UpperCamelCase       				=model(A_   ,  labels=A_       )
         self.parent.assertEqual(result.logits.shape   ,  (self.batch_size, self.num_labels)       )

         # test greyscale images
         __UpperCamelCase       				=1
         __UpperCamelCase       				=DinatForImageClassification(A_       )
         model.to(A_       )
         model.eval()

         __UpperCamelCase       				=floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
         __UpperCamelCase       				=model(A_       )
         self.parent.assertEqual(result.logits.shape   ,  (self.batch_size, self.num_labels)       )






    def _a   (     self   ,  A_   ,  A_   ,  A_       )    ->			int:
         __UpperCamelCase       				=DinatBackbone(config=A_       )
         model.to(A_       )
         model.eval()
         __UpperCamelCase       				=model(A_       )

         # verify hidden states
         self.parent.assertEqual(len(result.feature_maps       )   ,  len(config.out_features       )       )
         self.parent.assertListEqual(list(result.feature_maps[0].shape       )   ,  [self.batch_size, model.channels[0], 16, 16]       )

         # verify channels
         self.parent.assertEqual(len(model.channels       )   ,  len(config.out_features       )       )

         # verify backbone works with out_features=None
         __UpperCamelCase       				=None
         __UpperCamelCase       				=DinatBackbone(config=A_       )
         model.to(A_       )
         model.eval()
         __UpperCamelCase       				=model(A_       )

         # verify feature maps
         self.parent.assertEqual(len(result.feature_maps       )   ,  1       )
         self.parent.assertListEqual(list(result.feature_maps[0].shape       )   ,  [self.batch_size, model.channels[-1], 4, 4]       )

         # verify channels
         self.parent.assertEqual(len(model.channels       )   ,  1       )






    def _a   (     self       )    ->			str:
         __UpperCamelCase       				=self.prepare_config_and_inputs()
         __UpperCamelCase				, __UpperCamelCase				, __UpperCamelCase       				=config_and_inputs
         __UpperCamelCase       				={'pixel_values': pixel_values}
         return config, inputs_dict






@require_natten
@require_torch
class        UpperCAmelCase__			(   A_  ,    A_  ,    unittest.TestCase						):



    """simple docstring"""






    UpperCAmelCase__		:						Dict    =    (
        (
            DinatModel,
            DinatForImageClassification,
            DinatBackbone,
        )
        if is_torch_available()
        else ()
    )
    UpperCAmelCase__		:						Dict    =    (
        {"feature-extraction": DinatModel, "image-classification": DinatForImageClassification}
        if is_torch_available()
        else {}
    )
    UpperCAmelCase__		:						List[Any]    =    False

    UpperCAmelCase__		:						str    =    False
    UpperCAmelCase__		:						str    =    False
    UpperCAmelCase__		:						List[str]    =    False
    UpperCAmelCase__		:						int    =    False






    def _a   (     self       )    ->			Optional[int]:
         __UpperCamelCase       				=DinatModelTester(self       )
         __UpperCamelCase       				=ConfigTester(self   ,  config_class=A_   ,  embed_dim=37       )






    def _a   (     self       )    ->			Optional[Any]:
         self.create_and_test_config_common_properties()
         self.config_tester.create_and_test_config_to_json_string()
         self.config_tester.create_and_test_config_to_json_file()
         self.config_tester.create_and_test_config_from_and_save_pretrained()
         self.config_tester.create_and_test_config_with_num_labels()
         self.config_tester.check_config_can_be_init_without_params()
         self.config_tester.check_config_arguments_init()






    def _a   (     self       )    ->			Optional[int]:
         return






    def _a   (     self       )    ->			List[Any]:
         __UpperCamelCase       				=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*A_       )






    def _a   (     self       )    ->			Dict:
         __UpperCamelCase       				=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*A_       )






    def _a   (     self       )    ->			str:
         __UpperCamelCase       				=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_backbone(*A_       )






    @unittest.skip(reason='Dinat does not use inputs_embeds'       )
    def _a   (     self       )    ->			Tuple:
         pass






    @unittest.skip(reason='Dinat does not use feedforward chunking'       )
    def _a   (     self       )    ->			Optional[Any]:
         pass






    def _a   (     self       )    ->			Union[str, Any]:
         __UpperCamelCase				, __UpperCamelCase       				=self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
              __UpperCamelCase       				=model_class(A_       )
              self.assertIsInstance(model.get_input_embeddings()   ,  (nn.Module)       )
              __UpperCamelCase       				=model.get_output_embeddings()
              self.assertTrue(x is None or isinstance(A_   ,  nn.Linear       )       )






    def _a   (     self       )    ->			Optional[Any]:
         __UpperCamelCase				, __UpperCamelCase       				=self.model_tester.prepare_config_and_inputs_for_common()

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

              __UpperCamelCase       				=['pixel_values']
              self.assertListEqual(arg_names[:1]   ,  A_       )






    def _a   (     self       )    ->			List[str]:
         self.skipTest('Dinat\'s attention operation is handled entirely by NATTEN.'       )






    def _a   (     self   ,  A_   ,  A_   ,  A_   ,  A_       )    ->			List[str]:
         __UpperCamelCase       				=model_class(A_       )
         model.to(A_       )
         model.eval()

         with torch.no_grad():
              __UpperCamelCase       				=model(**self._prepare_for_class(A_   ,  A_       )       )

         __UpperCamelCase       				=outputs.hidden_states

         __UpperCamelCase       				=getattr(
             self.model_tester   ,  'expected_num_hidden_layers'   ,  len(self.model_tester.depths       ) + 1       )
         self.assertEqual(len(A_       )   ,  A_       )

         # Dinat has a different seq_length
         __UpperCamelCase       				=(
             config.patch_size
             if isinstance(config.patch_size   ,  collections.abc.Iterable       )
             else (config.patch_size, config.patch_size)
         )

         __UpperCamelCase       				=image_size[0] // patch_size[0]
         __UpperCamelCase       				=image_size[1] // patch_size[1]

         self.assertListEqual(
             list(hidden_states[0].shape[-3:]       )   ,  [height, width, self.model_tester.embed_dim]   ,  )

         if model_class.__name__ != "DinatBackbone":
              __UpperCamelCase       				=outputs.reshaped_hidden_states
              self.assertEqual(len(A_       )   ,  A_       )

              __UpperCamelCase				, __UpperCamelCase				, __UpperCamelCase				, __UpperCamelCase       				=reshaped_hidden_states[0].shape
              __UpperCamelCase       				=(
                  reshaped_hidden_states[0].view(A_   ,  A_   ,  A_   ,  A_       ).permute(0   ,  2   ,  3   ,  1       )
              )
              self.assertListEqual(
                  list(reshaped_hidden_states.shape[-3:]       )   ,  [height, width, self.model_tester.embed_dim]   ,  )






    def _a   (     self       )    ->			Tuple:
         __UpperCamelCase				, __UpperCamelCase       				=self.model_tester.prepare_config_and_inputs_for_common()

         __UpperCamelCase       				=(
             self.model_tester.image_size
             if isinstance(self.model_tester.image_size   ,  collections.abc.Iterable       )
             else (self.model_tester.image_size, self.model_tester.image_size)
         )

         for model_class in self.all_model_classes:
              __UpperCamelCase       				=True
              self.check_hidden_states_output(A_   ,  A_   ,  A_   ,  A_       )

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              __UpperCamelCase       				=True

              self.check_hidden_states_output(A_   ,  A_   ,  A_   ,  A_       )






    @slow
    def _a   (     self       )    ->			Any:
         for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              __UpperCamelCase       				=DinatModel.from_pretrained(A_       )
              self.assertIsNotNone(A_       )






    def _a   (     self       )    ->			List[Any]:
         __UpperCamelCase				, __UpperCamelCase       				=self.model_tester.prepare_config_and_inputs_for_common()

         __UpperCamelCase       				=_config_zero_init(A_       )
         for model_class in self.all_model_classes:
              __UpperCamelCase       				=model_class(config=A_       )
              for name, param in model.named_parameters():
                   if "embeddings" not in name and param.requires_grad:
                        self.assertIn(
                            ((param.data.mean() * 1E9).round() / 1E9).item()   ,  [0.0, 1.0]   ,  msg=f'Parameter {name} of model {model_class} seems not properly initialized'   ,  )






@require_natten
@require_vision
@require_torch
class        UpperCAmelCase__			(   unittest.TestCase						):



    """simple docstring"""






    @cached_property
    def _a   (     self       )    ->			List[str]:
         return AutoImageProcessor.from_pretrained('shi-labs/dinat-mini-in1k-224'       ) if is_vision_available() else None






    @slow
    def _a   (     self       )    ->			Tuple:
         __UpperCamelCase       				=DinatForImageClassification.from_pretrained('shi-labs/dinat-mini-in1k-224'       ).to(A_       )
         __UpperCamelCase       				=self.default_image_processor

         __UpperCamelCase       				=Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'       )
         __UpperCamelCase       				=image_processor(images=A_   ,  return_tensors='pt'       ).to(A_       )

         # forward pass
         with torch.no_grad():
              __UpperCamelCase       				=model(**A_       )

         # verify the logits
         __UpperCamelCase       				=torch.Size((1, 1000)       )
         self.assertEqual(outputs.logits.shape   ,  A_       )
         __UpperCamelCase       				=torch.tensor([-0.1545, -0.7667, 0.4642]       ).to(A_       )
         self.assertTrue(torch.allclose(outputs.logits[0, :3]   ,  A_   ,  atol=1E-4       )       )






@require_torch
@require_natten
class        UpperCAmelCase__			(   unittest.TestCase  ,    A_						):



    """simple docstring"""






    UpperCAmelCase__		:						List[str]    =    (DinatBackbone,) if is_torch_available() else ()
    UpperCAmelCase__		:						List[Any]    =    DinatConfig






    def _a   (     self       )    ->			int:
         __UpperCamelCase       				=DinatModelTester(self       )









def 		_UpperCAmelCase							(     SCREAMING_SNAKE_CASE__     :   list     ,	SCREAMING_SNAKE_CASE__     :   int     ,	SCREAMING_SNAKE_CASE__     :   int     ,	SCREAMING_SNAKE_CASE__     :   int							):
     __UpperCamelCase       				=int((lower_bound + upper_bound) // 2							)
     if array[r] == value:
          return r
     if lower_bound >= upper_bound:
          return -1
     if array[r] < value:
          return binary_search(SCREAMING_SNAKE_CASE__     ,	r + 1     ,	SCREAMING_SNAKE_CASE__     ,	SCREAMING_SNAKE_CASE__							)
     else:
          return binary_search(SCREAMING_SNAKE_CASE__     ,	SCREAMING_SNAKE_CASE__     ,	r - 1     ,	SCREAMING_SNAKE_CASE__							)

def 		_UpperCAmelCase							(     SCREAMING_SNAKE_CASE__     :   int     ,	SCREAMING_SNAKE_CASE__     :   list							):
     __UpperCamelCase       				=0
     if matrix[index][0] == value:
          return [index, 0]
     while index < len(SCREAMING_SNAKE_CASE__							) and matrix[index][0] < value:
          __UpperCamelCase       				=binary_search(matrix[index]     ,	0     ,	len(matrix[index]							) - 1     ,	SCREAMING_SNAKE_CASE__							)
          if r != -1:
               return [index, r]
          index += 1
     return [-1, -1]


if __name__ == "__main__":
   import doctest

   doctest.testmod()





