




"""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  __lowercase    :


				'''simple docstring'''


				def __init__(				self				:       List[str] ,SCREAMING_SNAKE_CASE_				:       Optional[int] ,SCREAMING_SNAKE_CASE_				:       List[str]=1_3 ,SCREAMING_SNAKE_CASE_				:       Dict=6_4 ,SCREAMING_SNAKE_CASE_				:       Any=4 ,SCREAMING_SNAKE_CASE_				:       Optional[int]=3 ,SCREAMING_SNAKE_CASE_				:       str=1_6 ,SCREAMING_SNAKE_CASE_				:       List[Any]=[1, 2, 1] ,SCREAMING_SNAKE_CASE_				:       str=[2, 4, 8] ,SCREAMING_SNAKE_CASE_				:       Union[str, Any]=3 ,SCREAMING_SNAKE_CASE_				:       Optional[int]=2.0 ,SCREAMING_SNAKE_CASE_				:       str=True ,SCREAMING_SNAKE_CASE_				:       Union[str, Any]=0.0 ,SCREAMING_SNAKE_CASE_				:       Optional[Any]=0.0 ,SCREAMING_SNAKE_CASE_				:       int=0.1 ,SCREAMING_SNAKE_CASE_				:       Dict="gelu" ,SCREAMING_SNAKE_CASE_				:       Dict=True ,SCREAMING_SNAKE_CASE_				:       List[Any]=0.02 ,SCREAMING_SNAKE_CASE_				:       int=1E-5 ,SCREAMING_SNAKE_CASE_				:       Dict=True ,SCREAMING_SNAKE_CASE_				:       Any=None ,SCREAMING_SNAKE_CASE_				:       Optional[int]=True ,SCREAMING_SNAKE_CASE_				:       List[str]=1_0 ,SCREAMING_SNAKE_CASE_				:       Union[str, Any]=["stage1", "stage2"] ,SCREAMING_SNAKE_CASE_				:       Dict=[1, 2] ,):
										snake_case__    :   int							     =   parent
										snake_case__    :   Optional[int]							     =   batch_size
										snake_case__    :   int							     =   image_size
										snake_case__    :   Dict							     =   patch_size
										snake_case__    :   str							     =   num_channels
										snake_case__    :   Dict							     =   embed_dim
										snake_case__    :   List[Any]							     =   depths
										snake_case__    :   int							     =   num_heads
										snake_case__    :   Union[str, Any]							     =   kernel_size
										snake_case__    :   str							     =   mlp_ratio
										snake_case__    :   str							     =   qkv_bias
										snake_case__    :   Optional[int]							     =   hidden_dropout_prob
										snake_case__    :   Union[str, Any]							     =   attention_probs_dropout_prob
										snake_case__    :   Optional[Any]							     =   drop_path_rate
										snake_case__    :   List[str]							     =   hidden_act
										snake_case__    :   Dict							     =   patch_norm
										snake_case__    :   List[Any]							     =   layer_norm_eps
										snake_case__    :   Any							     =   initializer_range
										snake_case__    :   List[Any]							     =   is_training
										snake_case__    :   Optional[Any]							     =   scope
										snake_case__    :   Union[str, Any]							     =   use_labels
										snake_case__    :   Dict							     =   num_labels
										snake_case__    :   List[str]							     =   out_features
										snake_case__    :   List[str]							     =   out_indices

				def  _lowerCamelCase   (				self				:       List[Any] ):
										snake_case__    :   Any							     =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

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

										snake_case__    :   Union[str, Any]							     =   self.get_config()

										return config, pixel_values, labels

				def  _lowerCamelCase   (				self				:       Tuple ):
										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  _lowerCamelCase   (				self				:       Union[str, Any] ,SCREAMING_SNAKE_CASE_				:       Any ,SCREAMING_SNAKE_CASE_				:       str ,SCREAMING_SNAKE_CASE_				:       Union[str, Any] ):
										snake_case__    :   str							     =   NatModel(config=SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )
										model.eval()
										snake_case__    :   Union[str, Any]							     =   model(SCREAMING_SNAKE_CASE_ )

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

										# test greyscale images
										snake_case__    :   List[Any]							     =   1
										snake_case__    :   Optional[int]							     =   NatForImageClassification(SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )
										model.eval()

										snake_case__    :   int							     =   floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
										snake_case__    :   List[Any]							     =   model(SCREAMING_SNAKE_CASE_ )
										self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.num_labels) )

				def  _lowerCamelCase   (				self				:       List[Any] ,SCREAMING_SNAKE_CASE_				:       Tuple ,SCREAMING_SNAKE_CASE_				:       Dict ,SCREAMING_SNAKE_CASE_				:       Optional[int] ):
										snake_case__    :   Optional[int]							     =   NatBackbone(config=SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )
										model.eval()
										snake_case__    :   Optional[Any]							     =   model(SCREAMING_SNAKE_CASE_ )

										# 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
										snake_case__    :   Optional[Any]							     =   None
										snake_case__    :   Optional[Any]							     =   NatBackbone(config=SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )
										model.eval()
										snake_case__    :   int							     =   model(SCREAMING_SNAKE_CASE_ )

										# 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  _lowerCamelCase   (				self				:       List[str] ):
										snake_case__    :   Optional[int]							     =   self.prepare_config_and_inputs()
										snake_case__     ,		snake_case__     ,		snake_case__    :   Dict							     =   config_and_inputs
										snake_case__    :   List[str]							     =   {'''pixel_values''': pixel_values}
										return config, inputs_dict







@require_natten
@require_torch
class  __lowercase    (     snake_case_  ,				snake_case_  ,				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  _lowerCamelCase   (				self				:       Dict ):
										snake_case__    :   Optional[Any]							     =   NatModelTester(self )
										snake_case__    :   str							     =   ConfigTester(self ,config_class=SCREAMING_SNAKE_CASE_ ,embed_dim=3_7 )

				def  _lowerCamelCase   (				self				:       Optional[int] ):
										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  _lowerCamelCase   (				self				:       Union[str, Any] ):
										return

				def  _lowerCamelCase   (				self				:       Any ):
										snake_case__    :   Optional[int]							     =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_ )

				def  _lowerCamelCase   (				self				:       Any ):
										snake_case__    :   List[str]							     =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE_ )

				def  _lowerCamelCase   (				self				:       Any ):
										snake_case__    :   Tuple							     =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_backbone(*SCREAMING_SNAKE_CASE_ )

				@unittest.skip(reason='''Nat does not use inputs_embeds''' )
				def  _lowerCamelCase   (				self				:       List[Any] ):
										pass

				@unittest.skip(reason='''Nat does not use feedforward chunking''' )
				def  _lowerCamelCase   (				self				:       Union[str, Any] ):
										pass

				def  _lowerCamelCase   (				self				:       List[str] ):
										snake_case__     ,		snake_case__    :   str							     =   self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																snake_case__    :   List[Any]							     =   model_class(SCREAMING_SNAKE_CASE_ )
																self.assertIsInstance(model.get_input_embeddings() ,(nn.Module) )
																snake_case__    :   Dict							     =   model.get_output_embeddings()
																self.assertTrue(x is None or isinstance(SCREAMING_SNAKE_CASE_ ,nn.Linear ) )

				def  _lowerCamelCase   (				self				:       Dict ):
										snake_case__     ,		snake_case__    :   List[Any]							     =   self.model_tester.prepare_config_and_inputs_for_common()

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

																snake_case__    :   Optional[Any]							     =   ['''pixel_values''']
																self.assertListEqual(arg_names[:1] ,SCREAMING_SNAKE_CASE_ )

				def  _lowerCamelCase   (				self				:       Tuple ):
										self.skipTest('''Nat\'s attention operation is handled entirely by NATTEN.''' )

				def  _lowerCamelCase   (				self				:       Dict ,SCREAMING_SNAKE_CASE_				:       str ,SCREAMING_SNAKE_CASE_				:       List[str] ,SCREAMING_SNAKE_CASE_				:       List[Any] ,SCREAMING_SNAKE_CASE_				:       List[str] ):
										snake_case__    :   Optional[int]							     =   model_class(SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )
										model.eval()

										with torch.no_grad():
																snake_case__    :   Optional[int]							     =   model(**self._prepare_for_class(SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ ) )

										snake_case__    :   Dict							     =   outputs.hidden_states

										snake_case__    :   List[str]							     =   getattr(
										    self.model_tester ,'''expected_num_hidden_layers''' ,len(self.model_tester.depths ) + 1 )
										self.assertEqual(len(SCREAMING_SNAKE_CASE_ ) ,SCREAMING_SNAKE_CASE_ )

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

										snake_case__    :   Union[str, Any]							     =   image_size[0] // patch_size[0]
										snake_case__    :   List[Any]							     =   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":
																snake_case__    :   Optional[Any]							     =   outputs.reshaped_hidden_states
																self.assertEqual(len(SCREAMING_SNAKE_CASE_ ) ,SCREAMING_SNAKE_CASE_ )

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

				def  _lowerCamelCase   (				self				:       Optional[int] ):
										snake_case__     ,		snake_case__    :   Tuple							     =   self.model_tester.prepare_config_and_inputs_for_common()

										snake_case__    :   Dict							     =   (
										    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:
																snake_case__    :   Union[str, Any]							     =   True
																self.check_hidden_states_output(SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ )

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

																self.check_hidden_states_output(SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ )

				@slow
				def  _lowerCamelCase   (				self				:       List[str] ):
										for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																snake_case__    :   Optional[int]							     =   NatModel.from_pretrained(SCREAMING_SNAKE_CASE_ )
																self.assertIsNotNone(SCREAMING_SNAKE_CASE_ )





				def  _lowerCamelCase   (				self				:       List[Any] ):
										snake_case__     ,		snake_case__    :   int							     =   self.model_tester.prepare_config_and_inputs_for_common()

										snake_case__    :   List[Any]							     =   _config_zero_init(SCREAMING_SNAKE_CASE_ )
										for model_class in self.all_model_classes:
																snake_case__    :   Any							     =   model_class(config=SCREAMING_SNAKE_CASE_ )
																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  __lowercase    (     unittest.TestCase		):


				'''simple docstring'''


				@cached_property
				def  _lowerCamelCase   (				self				:       Any ):
										return AutoImageProcessor.from_pretrained('''shi-labs/nat-mini-in1k-224''' ) if is_vision_available() else None





				@slow
				def  _lowerCamelCase   (				self				:       str ):
										snake_case__    :   Optional[Any]							     =   NatForImageClassification.from_pretrained('''shi-labs/nat-mini-in1k-224''' ).to(SCREAMING_SNAKE_CASE_ )
										snake_case__    :   Union[str, Any]							     =   self.default_image_processor

										snake_case__    :   str							     =   Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''' )
										snake_case__    :   Any							     =   image_processor(images=SCREAMING_SNAKE_CASE_ ,return_tensors='''pt''' ).to(SCREAMING_SNAKE_CASE_ )

										# forward pass
										with torch.no_grad():
																snake_case__    :   Dict							     =   model(**SCREAMING_SNAKE_CASE_ )

										# verify the logits
										snake_case__    :   Any							     =   torch.Size((1, 1_0_0_0) )
										self.assertEqual(outputs.logits.shape ,SCREAMING_SNAKE_CASE_ )
										snake_case__    :   str							     =   torch.tensor([0.3805, -0.8676, -0.3912] ).to(SCREAMING_SNAKE_CASE_ )
										self.assertTrue(torch.allclose(outputs.logits[0, :3] ,SCREAMING_SNAKE_CASE_ ,atol=1E-4 ) )







@require_torch
@require_natten
class  __lowercase    (     unittest.TestCase  ,				snake_case_		):


				'''simple docstring'''


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

				def  _lowerCamelCase   (				self				:       Optional[int] ):
										snake_case__    :   Dict							     =   NatModelTester(self )








"""simple docstring"""





def 						lowerCAmelCase							(	lowerCamelCase						):
						'''simple docstring'''


						return int(lowerCamelCase						) if x - int(lowerCamelCase						) <= 0 else int(lowerCamelCase						) + 1


if __name__ == "__main__":
							import doctest

							doctest.testmod()


