






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			,			lowercase_			,			lowercase_=1_3			,			lowercase_=6_4			,			lowercase_=4			,			lowercase_=3			,			lowercase_=1_6			,			lowercase_=[1, 2, 1]			,			lowercase_=[2, 4, 8]			,			lowercase_=3			,			lowercase_=2.0			,			lowercase_=True			,			lowercase_=0.0			,			lowercase_=0.0			,			lowercase_=0.1			,			lowercase_="gelu"			,			lowercase_=True			,			lowercase_=0.02			,			lowercase_=1e-5			,			lowercase_=True			,			lowercase_=None			,			lowercase_=True			,			lowercase_=1_0			,			lowercase_=["stage1", "stage2"]			,			lowercase_=[1, 2]			,			)					->					List[str]:
														'''simple docstring'''

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

							def      _lowercase       (		self   )					->					Dict:
														'''simple docstring'''

														lowerCAmelCase_										=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

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

														lowerCAmelCase_										=     self.get_config()

														return config, pixel_values, labels

							def      _lowercase       (		self   )					->					List[Any]:
														'''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      _lowercase       (		self			,			lowercase_			,			lowercase_			,			lowercase_   )					->					Tuple:
														'''simple docstring'''

														lowerCAmelCase_										=     NatModel(config=lowercase_   )
														model.to(lowercase_   )
														model.eval()
														lowerCAmelCase_										=     model(lowercase_   )

														lowerCAmelCase_										=     lowerCAmelCase_										=     (config.image_size // config.patch_size) // (2 ** (len(config.depths   ) - 1))
														lowerCAmelCase_										=     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      _lowercase       (		self			,			lowercase_			,			lowercase_			,			lowercase_   )					->					Optional[int]:
														'''simple docstring'''

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

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

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

							def      _lowercase       (		self			,			lowercase_			,			lowercase_			,			lowercase_   )					->					Union[str, Any]:
														'''simple docstring'''

														lowerCAmelCase_										=     NatBackbone(config=lowercase_   )
														model.to(lowercase_   )
														model.eval()
														lowerCAmelCase_										=     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], 1_6, 1_6]   )

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

														# verify backbone works with out_features=None
														lowerCAmelCase_										=     None
														lowerCAmelCase_										=     NatBackbone(config=lowercase_   )
														model.to(lowercase_   )
														model.eval()
														lowerCAmelCase_										=     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      _lowercase       (		self   )					->					Tuple:
														'''simple docstring'''

														lowerCAmelCase_										=     self.prepare_config_and_inputs()
														lowerCAmelCase_						,   lowerCAmelCase_						,   lowerCAmelCase_										=     config_and_inputs
														lowerCAmelCase_										=     {'pixel_values': pixel_values}
														return config, inputs_dict




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


							'''simple docstring'''






							__a:	int     =    (
							    (
							        NatModel,
							        NatForImageClassification,
							        NatBackbone,
							    )
							    if is_torch_available()
							    else ()
							)
							__a:	List[Any]     =    (
							    {'''feature-extraction''': NatModel, '''image-classification''': NatForImageClassification}
							    if is_torch_available()
							    else {}
							)
							__a:	Optional[Any]     =    False

							__a:	Dict     =    False
							__a:	int     =    False
							__a:	List[str]     =    False
							__a:	List[Any]     =    False

							def      _lowercase       (		self   )					->					Tuple:
														'''simple docstring'''

														lowerCAmelCase_										=     NatModelTester(self   )
														lowerCAmelCase_										=     ConfigTester(self			,			config_class=lowercase_			,			embed_dim=3_7   )

							def      _lowercase       (		self   )					->					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      _lowercase       (		self   )					->					int:
														'''simple docstring'''

														return

							def      _lowercase       (		self   )					->					Optional[Any]:
														'''simple docstring'''

														lowerCAmelCase_										=     self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*lowercase_   )

							def      _lowercase       (		self   )					->					Optional[int]:
														'''simple docstring'''

														lowerCAmelCase_										=     self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_image_classification(*lowercase_   )

							def      _lowercase       (		self   )					->					Optional[int]:
														'''simple docstring'''

														lowerCAmelCase_										=     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      _lowercase       (		self   )					->					Optional[Any]:
														'''simple docstring'''

														pass

							@unittest.skip(reason='Nat does not use feedforward chunking'   )
							def      _lowercase       (		self   )					->					Union[str, Any]:
														'''simple docstring'''

														pass

							def      _lowercase       (		self   )					->					Union[str, Any]:
														'''simple docstring'''

														lowerCAmelCase_						,   lowerCAmelCase_										=     self.model_tester.prepare_config_and_inputs_for_common()

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

							def      _lowercase       (		self   )					->					Optional[Any]:
														'''simple docstring'''

														lowerCAmelCase_						,   lowerCAmelCase_										=     self.model_tester.prepare_config_and_inputs_for_common()

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

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

							def      _lowercase       (		self   )					->					List[str]:
														'''simple docstring'''

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

							def      _lowercase       (		self			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_   )					->					Tuple:
														'''simple docstring'''

														lowerCAmelCase_										=     model_class(lowercase_   )
														model.to(lowercase_   )
														model.eval()

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

														lowerCAmelCase_										=     outputs.hidden_states

														lowerCAmelCase_										=     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
														lowerCAmelCase_										=     (
														    config.patch_size
														    if isinstance(config.patch_size			,			collections.abc.Iterable   )
														    else (config.patch_size, config.patch_size)
														)

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

																					lowerCAmelCase_						,   lowerCAmelCase_						,   lowerCAmelCase_						,   lowerCAmelCase_										=     reshaped_hidden_states[0].shape
																					lowerCAmelCase_										=     (
																					    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      _lowercase       (		self   )					->					List[str]:
														'''simple docstring'''

														lowerCAmelCase_						,   lowerCAmelCase_										=     self.model_tester.prepare_config_and_inputs_for_common()

														lowerCAmelCase_										=     (
														    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:
																					lowerCAmelCase_										=     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"]
																					lowerCAmelCase_										=     True

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

							@slow
							def      _lowercase       (		self   )					->					Optional[int]:
														'''simple docstring'''

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






							def      _lowercase       (		self   )					->					Any:
														'''simple docstring'''

														lowerCAmelCase_						,   lowerCAmelCase_										=     self.model_tester.prepare_config_and_inputs_for_common()

														lowerCAmelCase_										=     _config_zero_init(lowercase_   )
														for model_class in self.all_model_classes:
																					lowerCAmelCase_										=     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      _lowercase       (		self   )					->					Any:
														'''simple docstring'''

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






							@slow
							def      _lowercase       (		self   )					->					Dict:
														'''simple docstring'''

														lowerCAmelCase_										=     NatForImageClassification.from_pretrained('shi-labs/nat-mini-in1k-224'   ).to(lowercase_   )
														lowerCAmelCase_										=     self.default_image_processor

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

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

														# verify the logits
														lowerCAmelCase_										=     torch.Size((1, 1_0_0_0)   )
														self.assertEqual(outputs.logits.shape			,			lowercase_   )
														lowerCAmelCase_										=     torch.tensor([0.38_05, -0.86_76, -0.39_12]   ).to(lowercase_   )
														self.assertTrue(torch.allclose(outputs.logits[0, :3]			,			lowercase_			,			atol=1e-4   )   )




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


							'''simple docstring'''






							__a:	List[str]     =    (NatBackbone,) if is_torch_available() else ()
							__a:	str     =    NatConfig

							def      _lowercase       (		self   )					->					List[Any]:
														'''simple docstring'''

														lowerCAmelCase_										=     NatModelTester(self   )











import math
from collections import OrderedDict

import torch
from packaging import version
from torch import Tensor, nn

from .utils import logging


lowerCamelCase_						     =  logging.get_logger(__name__)




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self   )					->					Any:
														'''simple docstring'''

														super().__init__()
														if version.parse(torch.__version__   ) < version.parse('1.12.0'   ):
																					raise ImportError(
																					    f'''You are using torch=={torch.__version__}, but torch>=1.12.0 is required to use '''
																					    'PytorchGELUTanh. Please upgrade torch.'   )






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return nn.functional.gelu(lowercase_			,			approximate='tanh'   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi   ) * (input + 0.04_47_15 * torch.pow(lowercase_			,			3.0   ))   ))




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self			,			lowercase_ = False   )					->					Optional[Any]:
														'''simple docstring'''

														super().__init__()
														if use_gelu_python:
																					lowerCAmelCase_										=     self._gelu_python
														else:
																					lowerCAmelCase_										=     nn.functional.gelu

							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0   )   ))






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return self.act(lowercase_   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return 0.5 * input * (1.0 + torch.tanh(input * 0.79_78_84_56_08 * (1.0 + 0.04_47_15 * input * input)   ))




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return input * torch.sigmoid(1.7_02 * input   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self			,			lowercase_			,			lowercase_   )					->					Optional[int]:
														'''simple docstring'''

														if min > max:
																					raise ValueError(f'''min should be < max (got min: {min}, max: {max})'''   )

														super().__init__()
														lowerCAmelCase_										=     min
														lowerCAmelCase_										=     max






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return torch.clip(gelu(lowercase_   )			,			self.min			,			self.max   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self   )					->					Union[str, Any]:
														'''simple docstring'''

														super().__init__()
														lowerCAmelCase_										=     math.sqrt(2 / math.pi   )






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return 0.5 * input * (1 + torch.tanh(self.precomputed_constant * (input + 0.04_47_15 * torch.pow(lowercase_			,			3   ))   ))




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return nn.functional.silu(lowercase_   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def __init__(		self   )					->					List[Any]:
														'''simple docstring'''

														super().__init__()
														if version.parse(torch.__version__   ) < version.parse('1.9.0'   ):
																					lowerCAmelCase_										=     self._mish_python
														else:
																					lowerCAmelCase_										=     nn.functional.mish

							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return input * torch.tanh(nn.functional.softplus(lowercase_   )   )






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return self.act(lowercase_   )




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Tensor:
														'''simple docstring'''

														return input




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_			,			lowercase_=0.70_71_07			,			lowercase_=0.28_20_95   )					->					Optional[Any]:
														'''simple docstring'''

														lowerCAmelCase_										=     (input - mu).div(sigma * math.sqrt(2.0   )   )
														return 0.5 * (1.0 + torch.erf(lowercase_   ))




class 							a_ (						nn.Module		):


							'''simple docstring'''






							def      _lowercase       (		self			,			lowercase_   )					->					Any:
														'''simple docstring'''

														lowerCAmelCase_										=     nn.functional.relu(lowercase_   )
														lowerCAmelCase_										=     torch.square(lowercase_   )
														return squared




class 							a_ (						a_		):


							'''simple docstring'''






							def __getitem__(		self			,			lowercase_   )					->					Any:
														'''simple docstring'''

														lowerCAmelCase_										=     super().__getitem__(lowercase_   )
														lowerCAmelCase_						,   lowerCAmelCase_										=     content if isinstance(lowercase_			,			lowercase_   ) else (content, {})
														return cls(**lowercase_   )


lowerCamelCase_						     =  {
    """gelu""": GELUActivation,
    """gelu_10""": (ClippedGELUActivation, {"""min""": -1_0, """max""": 1_0}),
    """gelu_fast""": FastGELUActivation,
    """gelu_new""": NewGELUActivation,
    """gelu_python""": (GELUActivation, {"""use_gelu_python""": True}),
    """gelu_pytorch_tanh""": PytorchGELUTanh,
    """gelu_accurate""": AccurateGELUActivation,
    """laplace""": LaplaceActivation,
    """linear""": LinearActivation,
    """mish""": MishActivation,
    """quick_gelu""": QuickGELUActivation,
    """relu""": nn.ReLU,
    """relu2""": ReLUSquaredActivation,
    """relu6""": nn.ReLUa,
    """sigmoid""": nn.Sigmoid,
    """silu""": SiLUActivation,
    """swish""": SiLUActivation,
    """tanh""": nn.Tanh,
}
lowerCamelCase_						     =  ClassInstantier(ACTaCLS)
def      lowerCamelCase					(						a_	)	->			str:
							if activation_string in ACTaFN:
														return ACTaFN[activation_string]
							else:
														raise KeyError(F'''function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()	)}'''	)


# For backwards compatibility with: from activations import gelu_python
lowerCamelCase_						     =  get_activation("""gelu_python""")
lowerCamelCase_						     =  get_activation("""gelu_new""")
lowerCamelCase_						     =  get_activation("""gelu""")
lowerCamelCase_						     =  get_activation("""gelu_fast""")
lowerCamelCase_						     =  get_activation("""quick_gelu""")
lowerCamelCase_						     =  get_activation("""silu""")
lowerCamelCase_						     =  get_activation("""mish""")



lowerCamelCase_						     =  get_activation("""linear""")



