




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







			def __init__(	self   :      str   ,   snake_case__   :      Optional[Any]   ,   snake_case__   :      Tuple=13   ,   snake_case__   :      Optional[Any]=64   ,   snake_case__   :      Union[str, Any]=4   ,   snake_case__   :      str=3   ,   snake_case__   :      int=16   ,   snake_case__   :      int=[1, 2, 1]   ,   snake_case__   :      Optional[Any]=[2, 4, 8]   ,   snake_case__   :      List[Any]=3   ,   snake_case__   :      List[str]=2.0   ,   snake_case__   :      Optional[int]=True   ,   snake_case__   :      Dict=0.0   ,   snake_case__   :      List[Any]=0.0   ,   snake_case__   :      int=0.1   ,   snake_case__   :      Any="gelu"   ,   snake_case__   :      Optional[int]=True   ,   snake_case__   :      Any=0.02   ,   snake_case__   :      Tuple=1e-5   ,   snake_case__   :      Tuple=True   ,   snake_case__   :      Dict=None   ,   snake_case__   :      Dict=True   ,   snake_case__   :      List[str]=10   ,   snake_case__   :      str=["stage1", "stage2"]   ,   snake_case__   :      Tuple=[1, 2]   ,   ):




									"""simple docstring"""


									_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													=					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 UpperCamelCase      (	self   :      Any				):




									"""simple docstring"""


									_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 UpperCamelCase      (	self   :      str				):




									"""simple docstring"""


									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 UpperCamelCase      (	self   :      Union[str, Any]   ,   snake_case__   :      Optional[Any]   ,   snake_case__   :      Dict   ,   snake_case__   :      List[Any]				):




									"""simple docstring"""


									_UpperCAmelCase													=					NatModel(config=snake_case__				)
									model.to(snake_case__				)
									model.eval()
									_UpperCAmelCase													=					model(snake_case__				)

									_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 UpperCamelCase      (	self   :      Optional[Any]   ,   snake_case__   :      int   ,   snake_case__   :      Tuple   ,   snake_case__   :      Union[str, Any]				):




									"""simple docstring"""


									_UpperCAmelCase													=					NatForImageClassification(snake_case__				)
									model.to(snake_case__				)
									model.eval()
									_UpperCAmelCase													=					model(snake_case__   ,   labels=snake_case__				)
									self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.num_labels)				)

									# test greyscale images
									_UpperCAmelCase													=					1
									_UpperCAmelCase													=					NatForImageClassification(snake_case__				)
									model.to(snake_case__				)
									model.eval()

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







			def UpperCamelCase      (	self   :      Optional[Any]   ,   snake_case__   :      Optional[Any]   ,   snake_case__   :      Dict   ,   snake_case__   :      Optional[int]				):




									"""simple docstring"""


									_UpperCAmelCase													=					NatBackbone(config=snake_case__				)
									model.to(snake_case__				)
									model.eval()
									_UpperCAmelCase													=					model(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], 16, 16]				)

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

									# verify backbone works with out_features=None
									_UpperCAmelCase													=					None
									_UpperCAmelCase													=					NatBackbone(config=snake_case__				)
									model.to(snake_case__				)
									model.eval()
									_UpperCAmelCase													=					model(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 UpperCamelCase      (	self   :      List[Any]				):




									"""simple docstring"""


									_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    __lowerCAmelCase			(	UpperCAmelCase__ ,		UpperCAmelCase__ ,		unittest.TestCase			):
			snake_case_				:			Any								=					(
			    (
			        NatModel,
			        NatForImageClassification,
			        NatBackbone,
			    )
			    if is_torch_available()
			    else ()
			)
			snake_case_				:			str								=					(
			    {"feature-extraction": NatModel, "image-classification": NatForImageClassification}
			    if is_torch_available()
			    else {}
			)
			snake_case_				:			List[str]								=					False

			snake_case_				:			List[str]								=					False
			snake_case_				:			Tuple								=					False
			snake_case_				:			Dict								=					False
			snake_case_				:			Any								=					False







			def UpperCamelCase      (	self   :      Dict				):




									"""simple docstring"""


									_UpperCAmelCase													=					NatModelTester(self				)
									_UpperCAmelCase													=					ConfigTester(self   ,   config_class=snake_case__   ,   embed_dim=37				)







			def UpperCamelCase      (	self   :      Any				):




									"""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   :      Union[str, Any]				):




									"""simple docstring"""


									return







			def UpperCamelCase      (	self   :      Union[str, Any]				):




									"""simple docstring"""


									_UpperCAmelCase													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*snake_case__				)







			def UpperCamelCase      (	self   :      int				):




									"""simple docstring"""


									_UpperCAmelCase													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*snake_case__				)







			def UpperCamelCase      (	self   :      Optional[Any]				):




									"""simple docstring"""


									_UpperCAmelCase													=					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_backbone(*snake_case__				)







			@unittest.skip(reason="Nat does not use inputs_embeds"				)
			def UpperCamelCase      (	self   :      str				):




									"""simple docstring"""


									pass







			@unittest.skip(reason="Nat does not use feedforward chunking"				)
			def UpperCamelCase      (	self   :      int				):




									"""simple docstring"""


									pass







			def UpperCamelCase      (	self   :      int				):




									"""simple docstring"""


									_UpperCAmelCase							,			_UpperCAmelCase													=					self.model_tester.prepare_config_and_inputs_for_common()

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







			def UpperCamelCase      (	self   :      Dict				):




									"""simple docstring"""


									_UpperCAmelCase							,			_UpperCAmelCase													=					self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															_UpperCAmelCase													=					model_class(snake_case__				)
															_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]   ,   snake_case__				)







			def UpperCamelCase      (	self   :      Union[str, Any]				):




									"""simple docstring"""


									self.skipTest("Nat's attention operation is handled entirely by NATTEN."				)







			def UpperCamelCase      (	self   :      Optional[Any]   ,   snake_case__   :      Tuple   ,   snake_case__   :      Optional[Any]   ,   snake_case__   :      Union[str, Any]   ,   snake_case__   :      int				):




									"""simple docstring"""


									_UpperCAmelCase													=					model_class(snake_case__				)
									model.to(snake_case__				)
									model.eval()

									with torch.no_grad():
															_UpperCAmelCase													=					model(**self._prepare_for_class(snake_case__   ,   snake_case__				)				)

									_UpperCAmelCase													=					outputs.hidden_states

									_UpperCAmelCase													=					getattr(
									    self.model_tester   ,   "expected_num_hidden_layers"   ,   len(self.model_tester.depths				) + 1				)
									self.assertEqual(len(snake_case__				)   ,   snake_case__				)

									# Nat 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__ != "NatBackbone":
															_UpperCAmelCase													=					outputs.reshaped_hidden_states
															self.assertEqual(len(snake_case__				)   ,   snake_case__				)

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







			def UpperCamelCase      (	self   :      Optional[Any]				):




									"""simple docstring"""


									_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(snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__				)

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

															self.check_hidden_states_output(snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__				)







			@slow
			def UpperCamelCase      (	self   :      List[str]				):




									"""simple docstring"""


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







			def UpperCamelCase      (	self   :      Union[str, Any]				):




									"""simple docstring"""


									_UpperCAmelCase							,			_UpperCAmelCase													=					self.model_tester.prepare_config_and_inputs_for_common()

									_UpperCAmelCase													=					_config_zero_init(snake_case__				)
									for model_class in self.all_model_classes:
															_UpperCAmelCase													=					model_class(config=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    __lowerCAmelCase			(	unittest.TestCase			):







			@cached_property
			def UpperCamelCase      (	self   :      Optional[int]				):




									"""simple docstring"""


									return AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224"				) if is_vision_available() else None







			@slow
			def UpperCamelCase      (	self   :      List[str]				):




									"""simple docstring"""


									_UpperCAmelCase													=					NatForImageClassification.from_pretrained("shi-labs/nat-mini-in1k-224"				).to(snake_case__				)
									_UpperCAmelCase													=					self.default_image_processor

									_UpperCAmelCase													=					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"				)
									_UpperCAmelCase													=					image_processor(images=snake_case__   ,   return_tensors="pt"				).to(snake_case__				)

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

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






@require_torch
@require_natten
class    __lowerCAmelCase			(	unittest.TestCase ,		UpperCAmelCase__			):
			snake_case_				:			Optional[int]								=					(NatBackbone,) if is_torch_available() else ()
			snake_case_				:			Optional[int]								=					NatConfig







			def UpperCamelCase      (	self   :      Any				):




									"""simple docstring"""


									_UpperCAmelCase													=					NatModelTester(self				)










from __future__ import annotations


def 							__SCREAMING_SNAKE_CASE  (      snake_case_     ,   snake_case_						):

						'''simple docstring'''


						_UpperCAmelCase													=					{}
						for index, val in enumerate(snake_case_						):
												_UpperCAmelCase													=					target - val
												if compl in chk_map:
																		return [chk_map[compl], index]
												_UpperCAmelCase													=					index
						return []


if __name__ == "__main__":
				import doctest

				doctest.testmod()
				print(f"""{two_sum([2, 7, 11, 15], 9) = }""")




