


"""simple docstring"""







import collections
import inspect
import unittest

from transformers import NatConfig
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 NatBackbone, NatForImageClassification, NatModel
       from transformers.models.nat.modeling_nat import NAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
       from PIL import Image

       from transformers import AutoImageProcessor






class  __A							:






 """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=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:
     a											=parent
     a											=batch_size
     a											=image_size
     a											=patch_size
     a											=num_channels
     a											=embed_dim
     a											=depths
     a											=num_heads
     a											=kernel_size
     a											=mlp_ratio
     a											=qkv_bias
     a											=hidden_dropout_prob
     a											=attention_probs_dropout_prob
     a											=drop_path_rate
     a											=hidden_act
     a											=patch_norm
     a											=layer_norm_eps
     a											=initializer_range
     a											=is_training
     a											=scope
     a											=use_labels
     a											=num_labels
     a											=out_features
     a											=out_indices


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Optional[int]:
     a											=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

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

     a											=self.get_config()

     return config, pixel_values, labels


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	List[str]:
     return NatConfig(
         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					,       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  SCREAMING_SNAKE_CASE				(						self					,       __A					,       __A					,       __A		)   ->	List[str]:
     a											=NatModel(config=lowercase_		)
     model.to(lowercase_		)
     model.eval()
     a											=model(lowercase_		)

     a											=a											=(config.image_size // config.patch_size) // (2 ** (len(config.depths		) - 1))
     a											=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  SCREAMING_SNAKE_CASE				(						self					,       __A					,       __A					,       __A		)   ->	Union[str, Any]:
     a											=NatForImageClassification(lowercase_		)
     model.to(lowercase_		)
     model.eval()
     a											=model(lowercase_					,       labels=lowercase_		)
     self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.num_labels)		)

     # test greyscale images
     a											=1
     a											=NatForImageClassification(lowercase_		)
     model.to(lowercase_		)
     model.eval()

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


 def  SCREAMING_SNAKE_CASE				(						self					,       __A					,       __A					,       __A		)   ->	Union[str, Any]:
     a											=NatBackbone(config=lowercase_		)
     model.to(lowercase_		)
     model.eval()
     a											=model(lowercase_		)

     # 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
     a											=None
     a											=NatBackbone(config=lowercase_		)
     model.to(lowercase_		)
     model.eval()
     a											=model(lowercase_		)

     # 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  SCREAMING_SNAKE_CASE				(						self		)   ->	Dict:
     a											=self.prepare_config_and_inputs()
     a						, a						, a											=config_and_inputs
     a											={'''pixel_values''': pixel_values}
     return config, inputs_dict






@require_natten
@require_torch
class  __A							(			lowerCamelCase__,       lowerCamelCase__,       unittest.TestCase     ):






 """simple docstring"""




 __lowerCAmelCase					   =    (
     (
         NatModel,
         NatForImageClassification,
         NatBackbone,
     )
     if is_torch_available()
     else ()
 )
 __lowerCAmelCase					   =    (
     {'feature-extraction': NatModel, 'image-classification': NatForImageClassification}
     if is_torch_available()
     else {}
 )
 __lowerCAmelCase					   =    False

 __lowerCAmelCase					   =    False
 __lowerCAmelCase					   =    False
 __lowerCAmelCase					   =    False
 __lowerCAmelCase					   =    False


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	int:
     a											=NatModelTester(self		)
     a											=ConfigTester(self					,       config_class=lowercase_					,       embed_dim=37		)


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Dict:
     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  SCREAMING_SNAKE_CASE				(						self		)   ->	Tuple:
     return


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Any:
     a											=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_		)


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Optional[Any]:
     a											=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowercase_		)


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	int:
     a											=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_backbone(*lowercase_		)


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


 @unittest.skip(reason='''Nat does not use feedforward chunking'''		)
 def  SCREAMING_SNAKE_CASE				(						self		)   ->	List[str]:
     pass


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Tuple:
     a						, a											=self.model_tester.prepare_config_and_inputs_for_common()

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


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Optional[int]:
     a						, a											=self.model_tester.prepare_config_and_inputs_for_common()

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

         a											=['''pixel_values''']
         self.assertListEqual(arg_names[:1]					,       lowercase_		)


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Union[str, Any]:
     self.skipTest('''Nat\'s attention operation is handled entirely by NATTEN.'''		)


 def  SCREAMING_SNAKE_CASE				(						self					,       __A					,       __A					,       __A					,       __A		)   ->	List[Any]:
     a											=model_class(lowercase_		)
     model.to(lowercase_		)
     model.eval()

     with torch.no_grad():
         a											=model(**self._prepare_for_class(lowercase_					,       lowercase_		)		)

     a											=outputs.hidden_states

     a											=getattr(
         self.model_tester					,       '''expected_num_hidden_layers'''					,       len(self.model_tester.depths		) + 1		)
     self.assertEqual(len(lowercase_		)					,       lowercase_		)

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

     a											=image_size[0] // patch_size[0]
     a											=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__ != "NatBackbone":
         a											=outputs.reshaped_hidden_states
         self.assertEqual(len(lowercase_		)					,       lowercase_		)

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


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	int:
     a						, a											=self.model_tester.prepare_config_and_inputs_for_common()

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

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

         self.check_hidden_states_output(lowercase_					,       lowercase_					,       lowercase_					,       lowercase_		)


 @slow
 def  SCREAMING_SNAKE_CASE				(						self		)   ->	str:






     for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         a											=NatModel.from_pretrained(lowercase_		)
         self.assertIsNotNone(lowercase_		)





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

     a											=_config_zero_init(lowercase_		)
     for model_class in self.all_model_classes:
         a											=model_class(config=lowercase_		)
         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  SCREAMING_SNAKE_CASE				(						self		)   ->	List[str]:
     return AutoImageProcessor.from_pretrained('''shi-labs/nat-mini-in1k-224'''		) if is_vision_available() else None





 @slow
 def  SCREAMING_SNAKE_CASE				(						self		)   ->	List[str]:
     a											=NatForImageClassification.from_pretrained('''shi-labs/nat-mini-in1k-224'''		).to(lowercase_		)
     a											=self.default_image_processor

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

     # forward pass
     with torch.no_grad():
         a											=model(**lowercase_		)

     # verify the logits
     a											=torch.Size((1, 1000)		)
     self.assertEqual(outputs.logits.shape					,       lowercase_		)
     a											=torch.tensor([0.3_805, -0.8_676, -0.3_912]		).to(lowercase_		)
     self.assertTrue(torch.allclose(outputs.logits[0, :3]					,       lowercase_					,       atol=1E-4		)		)






@require_torch
@require_natten
class  __A							(			unittest.TestCase,       lowerCamelCase__     ):






 """simple docstring"""




 __lowerCAmelCase					   =    (NatBackbone,) if is_torch_available() else ()
 __lowerCAmelCase					   =    NatConfig


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Tuple:
     a											=NatModelTester(self		)


def      UpperCAmelCase__							(     _A :					int      , _A :					int				):





  '''simple docstring'''







  a__    =1.0
  for i in range(1      , r + 1				):
    ret *= (n + 1 - i) / i
  return round(_A				)




def      UpperCAmelCase__							(     _A :					int				):





  '''simple docstring'''







  return choose(8 + n      , _A				) + choose(9 + n      , _A				) - 10




def      UpperCAmelCase__							(     _A :					int				):





  '''simple docstring'''







  return sum(non_bouncy_exact(_A				) for i in range(1      , n + 1				)				)




def      UpperCAmelCase__							(     _A :					int = 1_00				):





  '''simple docstring'''







  return non_bouncy_upto(_A				)


if __name__ == "__main__":
       print(f"""{solution() = }""")
