






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 					a_:



     """simple docstring"""



     def __init__(						self			:					List[Any]       ,       lowerCAmelCase__			:					List[Any]       ,       lowerCAmelCase__			:					str=1_3       ,       lowerCAmelCase__			:					Dict=6_4       ,       lowerCAmelCase__			:					Union[str, Any]=4       ,       lowerCAmelCase__			:					int=3       ,       lowerCAmelCase__			:					int=1_6       ,       lowerCAmelCase__			:					Dict=[1, 2, 1]       ,       lowerCAmelCase__			:					int=[2, 4, 8]       ,       lowerCAmelCase__			:					Optional[int]=3       ,       lowerCAmelCase__			:					int=[[3], [1, 2], [1]]       ,       lowerCAmelCase__			:					Any=2.0       ,       lowerCAmelCase__			:					Optional[int]=True       ,       lowerCAmelCase__			:					str=0.0       ,       lowerCAmelCase__			:					Optional[int]=0.0       ,       lowerCAmelCase__			:					Optional[int]=0.1       ,       lowerCAmelCase__			:					List[Any]="gelu"       ,       lowerCAmelCase__			:					List[str]=True       ,       lowerCAmelCase__			:					int=0.02       ,       lowerCAmelCase__			:					List[str]=1e-5       ,       lowerCAmelCase__			:					Optional[int]=True       ,       lowerCAmelCase__			:					Optional[Any]=None       ,       lowerCAmelCase__			:					Union[str, Any]=True       ,       lowerCAmelCase__			:					int=1_0       ,       lowerCAmelCase__			:					str=["stage1", "stage2"]       ,       lowerCAmelCase__			:					Tuple=[1, 2]       ,       )    ->					Dict:

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						parent
      SCREAMING_SNAKE_CASE							  =						batch_size
      SCREAMING_SNAKE_CASE							  =						image_size
      SCREAMING_SNAKE_CASE							  =						patch_size
      SCREAMING_SNAKE_CASE							  =						num_channels
      SCREAMING_SNAKE_CASE							  =						embed_dim
      SCREAMING_SNAKE_CASE							  =						depths
      SCREAMING_SNAKE_CASE							  =						num_heads
      SCREAMING_SNAKE_CASE							  =						kernel_size
      SCREAMING_SNAKE_CASE							  =						dilations
      SCREAMING_SNAKE_CASE							  =						mlp_ratio
      SCREAMING_SNAKE_CASE							  =						qkv_bias
      SCREAMING_SNAKE_CASE							  =						hidden_dropout_prob
      SCREAMING_SNAKE_CASE							  =						attention_probs_dropout_prob
      SCREAMING_SNAKE_CASE							  =						drop_path_rate
      SCREAMING_SNAKE_CASE							  =						hidden_act
      SCREAMING_SNAKE_CASE							  =						patch_norm
      SCREAMING_SNAKE_CASE							  =						layer_norm_eps
      SCREAMING_SNAKE_CASE							  =						initializer_range
      SCREAMING_SNAKE_CASE							  =						is_training
      SCREAMING_SNAKE_CASE							  =						scope
      SCREAMING_SNAKE_CASE							  =						use_labels
      SCREAMING_SNAKE_CASE							  =						num_labels
      SCREAMING_SNAKE_CASE							  =						out_features
      SCREAMING_SNAKE_CASE							  =						out_indices





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

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

      SCREAMING_SNAKE_CASE							  =						self.get_config()

      return config, pixel_values, labels





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

      """simple docstring"""
      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   __UpperCamelCase		(						self			:					Dict       ,       lowerCAmelCase__			:					int       ,       lowerCAmelCase__			:					List[str]       ,       lowerCAmelCase__			:					Union[str, Any])    ->					Optional[int]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						DinatModel(config=lowerCAmelCase__)
      model.to(lowerCAmelCase__)
      model.eval()
      SCREAMING_SNAKE_CASE							  =						model(lowerCAmelCase__)

      SCREAMING_SNAKE_CASE							  =						SCREAMING_SNAKE_CASE							  =						(config.image_size // config.patch_size) // (2 ** (len(config.depths) - 1))
      SCREAMING_SNAKE_CASE							  =						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   __UpperCamelCase		(						self			:					str       ,       lowerCAmelCase__			:					Optional[Any]       ,       lowerCAmelCase__			:					Optional[Any]       ,       lowerCAmelCase__			:					Optional[Any])    ->					Dict:

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						DinatForImageClassification(lowerCAmelCase__)
      model.to(lowerCAmelCase__)
      model.eval()
      SCREAMING_SNAKE_CASE							  =						model(lowerCAmelCase__       ,       labels=lowerCAmelCase__)
      self.parent.assertEqual(result.logits.shape       ,       (self.batch_size, self.num_labels))

      # test greyscale images
      SCREAMING_SNAKE_CASE							  =						1
      SCREAMING_SNAKE_CASE							  =						DinatForImageClassification(lowerCAmelCase__)
      model.to(lowerCAmelCase__)
      model.eval()

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





     def   __UpperCamelCase		(						self			:					Dict       ,       lowerCAmelCase__			:					Optional[Any]       ,       lowerCAmelCase__			:					List[str]       ,       lowerCAmelCase__			:					List[Any])    ->					Tuple:

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						DinatBackbone(config=lowerCAmelCase__)
      model.to(lowerCAmelCase__)
      model.eval()
      SCREAMING_SNAKE_CASE							  =						model(lowerCAmelCase__)

      # 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], 1_6, 1_6])

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

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

      # 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   __UpperCamelCase		(						self			:					int)    ->					List[str]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						self.prepare_config_and_inputs()
      SCREAMING_SNAKE_CASE  , SCREAMING_SNAKE_CASE  , SCREAMING_SNAKE_CASE							  =						config_and_inputs
      SCREAMING_SNAKE_CASE							  =						{'pixel_values': pixel_values}
      return config, inputs_dict

@require_natten
@require_torch
class 					a_(						lowercase__	,    lowercase__	,    unittest.TestCase					):



     """simple docstring"""



     __snake_case						:							int       			=(
         (
             DinatModel,
             DinatForImageClassification,
             DinatBackbone,
         )
         if is_torch_available()
         else ()
     )
     __snake_case						:							str       			=(
         {'''feature-extraction''': DinatModel, '''image-classification''': DinatForImageClassification}
         if is_torch_available()
         else {}
     )
     __snake_case						:							Union[str, Any]       			=False

     __snake_case						:							List[Any]       			=False
     __snake_case						:							Tuple       			=False
     __snake_case						:							int       			=False
     __snake_case						:							List[str]       			=False





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

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						DinatModelTester(self)
      SCREAMING_SNAKE_CASE							  =						ConfigTester(self       ,       config_class=lowerCAmelCase__       ,       embed_dim=3_7)





     def   __UpperCamelCase		(						self			:					Tuple)    ->					Optional[int]:

      """simple docstring"""
      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   __UpperCamelCase		(						self			:					Optional[int])    ->					Union[str, Any]:

      """simple docstring"""
      return





     def   __UpperCamelCase		(						self			:					Tuple)    ->					List[str]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowerCAmelCase__)





     def   __UpperCamelCase		(						self			:					List[str])    ->					str:

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase__)





     def   __UpperCamelCase		(						self			:					Optional[Any])    ->					Dict:

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_backbone(*lowerCAmelCase__)





     @unittest.skip(reason='Dinat does not use inputs_embeds')
     def   __UpperCamelCase		(						self			:					int)    ->					Optional[Any]:

      """simple docstring"""
      pass





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

      """simple docstring"""
      pass





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

      """simple docstring"""
      SCREAMING_SNAKE_CASE  , SCREAMING_SNAKE_CASE							  =						self.model_tester.prepare_config_and_inputs_for_common()

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





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

      """simple docstring"""
      SCREAMING_SNAKE_CASE  , SCREAMING_SNAKE_CASE							  =						self.model_tester.prepare_config_and_inputs_for_common()

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

       SCREAMING_SNAKE_CASE							  =						['pixel_values']
       self.assertListEqual(arg_names[:1]       ,       lowerCAmelCase__)





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

      """simple docstring"""
      self.skipTest('Dinat\'s attention operation is handled entirely by NATTEN.')





     def   __UpperCamelCase		(						self			:					Optional[int]       ,       lowerCAmelCase__			:					Dict       ,       lowerCAmelCase__			:					List[str]       ,       lowerCAmelCase__			:					Union[str, Any]       ,       lowerCAmelCase__			:					str)    ->					Any:

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						model_class(lowerCAmelCase__)
      model.to(lowerCAmelCase__)
      model.eval()

      with torch.no_grad():
       SCREAMING_SNAKE_CASE							  =						model(**self._prepare_for_class(lowerCAmelCase__       ,       lowerCAmelCase__))

      SCREAMING_SNAKE_CASE							  =						outputs.hidden_states

      SCREAMING_SNAKE_CASE							  =						getattr(
          self.model_tester       ,       'expected_num_hidden_layers'       ,       len(self.model_tester.depths) + 1)
      self.assertEqual(len(lowerCAmelCase__)       ,       lowerCAmelCase__)

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

      SCREAMING_SNAKE_CASE							  =						image_size[0] // patch_size[0]
      SCREAMING_SNAKE_CASE							  =						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":
       SCREAMING_SNAKE_CASE							  =						outputs.reshaped_hidden_states
       self.assertEqual(len(lowerCAmelCase__)       ,       lowerCAmelCase__)

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





     def   __UpperCamelCase		(						self			:					Dict)    ->					int:

      """simple docstring"""
      SCREAMING_SNAKE_CASE  , SCREAMING_SNAKE_CASE							  =						self.model_tester.prepare_config_and_inputs_for_common()

      SCREAMING_SNAKE_CASE							  =						(
          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:
       SCREAMING_SNAKE_CASE							  =						True
       self.check_hidden_states_output(lowerCAmelCase__       ,       lowerCAmelCase__       ,       lowerCAmelCase__       ,       lowerCAmelCase__)

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

       self.check_hidden_states_output(lowerCAmelCase__       ,       lowerCAmelCase__       ,       lowerCAmelCase__       ,       lowerCAmelCase__)





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

      """simple docstring"""
      for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       SCREAMING_SNAKE_CASE							  =						DinatModel.from_pretrained(lowerCAmelCase__)
       self.assertIsNotNone(lowerCAmelCase__)





     def   __UpperCamelCase		(						self			:					Any)    ->					Any:

      """simple docstring"""
      SCREAMING_SNAKE_CASE  , SCREAMING_SNAKE_CASE							  =						self.model_tester.prepare_config_and_inputs_for_common()

      SCREAMING_SNAKE_CASE							  =						_config_zero_init(lowerCAmelCase__)
      for model_class in self.all_model_classes:
       SCREAMING_SNAKE_CASE							  =						model_class(config=lowerCAmelCase__)
       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 					a_(						unittest.TestCase					):



     """simple docstring"""



     @cached_property
     def   __UpperCamelCase		(						self			:					Optional[int])    ->					List[str]:

      """simple docstring"""
      return AutoImageProcessor.from_pretrained('shi-labs/dinat-mini-in1k-224') if is_vision_available() else None





     @slow
     def   __UpperCamelCase		(						self			:					Union[str, Any])    ->					Dict:

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						DinatForImageClassification.from_pretrained('shi-labs/dinat-mini-in1k-224').to(lowerCAmelCase__)
      SCREAMING_SNAKE_CASE							  =						self.default_image_processor

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

      # forward pass
      with torch.no_grad():
       SCREAMING_SNAKE_CASE							  =						model(**lowerCAmelCase__)

      # verify the logits
      SCREAMING_SNAKE_CASE							  =						torch.Size((1, 1_0_0_0))
      self.assertEqual(outputs.logits.shape       ,       lowerCAmelCase__)
      SCREAMING_SNAKE_CASE							  =						torch.tensor([-0.15_45, -0.76_67, 0.46_42]).to(lowerCAmelCase__)
      self.assertTrue(torch.allclose(outputs.logits[0, :3]       ,       lowerCAmelCase__       ,       atol=1e-4))

@require_torch
@require_natten
class 					a_(						unittest.TestCase	,    lowercase__					):



     """simple docstring"""



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





     def   __UpperCamelCase		(						self			:					Any)    ->					Tuple:

      """simple docstring"""
      SCREAMING_SNAKE_CASE							  =						DinatModelTester(self)












def 		A_			(		lowercase_	)							->int:







 """simple docstring"""

 if not isinstance(lowercase_					,       lowercase_	):
  SCREAMING_SNAKE_CASE							  =						f'''Input value of [number={number}] must be an integer'''
  raise TypeError(lowercase_	)
 if number < 1:
  raise ValueError('Input must be a positive integer'	)
 return number * (2 * number - 1)


if __name__ == "__main__":
      import doctest

      doctest.testmod()




