






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


   '''simple docstring'''

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

     '''simple docstring'''
     __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]  ):

     '''simple docstring'''
     __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			:					Tuple  ):

     '''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     _SCREAMING_SNAKE_CASE    (     self			:					Union[str, Any],	_lowerCamelCase			:					List[Any],	_lowerCamelCase			:					List[Any],	_lowerCamelCase			:					List[str]  ):

     '''simple docstring'''
     __A						      =			NatModel(config=_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()
     __A						      =			model(_lowerCamelCase  )

     __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			:					List[str],	_lowerCamelCase			:					Optional[Any],	_lowerCamelCase			:					Any,	_lowerCamelCase			:					List[str]  ):

     '''simple docstring'''
     __A						      =			NatForImageClassification(_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()
     __A						      =			model(_lowerCamelCase,	labels=_lowerCamelCase  )
     self.parent.assertEqual(result.logits.shape,	(self.batch_size, self.num_labels)  )

     # test greyscale images
     __A						      =			1
     __A						      =			NatForImageClassification(_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()

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


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[str],	_lowerCamelCase			:					Any,	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					int  ):

     '''simple docstring'''
     __A						      =			NatBackbone(config=_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()
     __A						      =			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], 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=_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()
     __A						      =			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     _SCREAMING_SNAKE_CASE    (     self			:					Tuple  ):

     '''simple docstring'''
     __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  snake_case		(							_lowerCAmelCase       ,     _lowerCAmelCase       ,     unittest.TestCase						):


   '''simple docstring'''

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

   A_						:     Any									=			False
   A_						:     Dict									=			False
   A_						:     str									=			False
   A_						:     Optional[Any]									=			False


   def     _SCREAMING_SNAKE_CASE    (     self			:					str  ):

     '''simple docstring'''
     __A						      =			NatModelTester(self  )
     __A						      =			ConfigTester(self,	config_class=_lowerCamelCase,	embed_dim=37  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[str]  ):

     '''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     _SCREAMING_SNAKE_CASE    (     self			:					Dict  ):

     '''simple docstring'''
     return


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[str]  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					str  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[Any]  ):

     '''simple docstring'''
     __A						      =			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_backbone(*_lowerCamelCase  )


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

     '''simple docstring'''
     pass


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

     '''simple docstring'''
     pass


   def     _SCREAMING_SNAKE_CASE    (     self			:					int  ):

     '''simple docstring'''
     __A      ,   __A						      =			self.model_tester.prepare_config_and_inputs_for_common()

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


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[str]  ):

     '''simple docstring'''
     __A      ,   __A						      =			self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       __A						      =			model_class(_lowerCamelCase  )
       __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],	_lowerCamelCase  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Any  ):

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


   def     _SCREAMING_SNAKE_CASE    (     self			:					Dict,	_lowerCamelCase			:					str,	_lowerCamelCase			:					Any,	_lowerCamelCase			:					Tuple,	_lowerCamelCase			:					Union[str, Any]  ):

     '''simple docstring'''
     __A						      =			model_class(_lowerCamelCase  )
     model.to(_lowerCamelCase  )
     model.eval()

     with torch.no_grad():
       __A						      =			model(**self._prepare_for_class(_lowerCamelCase,	_lowerCamelCase  )  )

     __A						      =			outputs.hidden_states

     __A						      =			getattr(
         self.model_tester,	'''expected_num_hidden_layers''',	len(self.model_tester.depths  ) + 1  )
     self.assertEqual(len(_lowerCamelCase  ),	_lowerCamelCase  )

     # 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(_lowerCamelCase  ),	_lowerCamelCase  )

       __A      ,   __A      ,   __A      ,   __A						      =			reshaped_hidden_states[0].shape
       __A						      =			(
           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     _SCREAMING_SNAKE_CASE    (     self			:					Union[str, Any]  ):

     '''simple docstring'''
     __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(_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase  )

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

       self.check_hidden_states_output(_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase  )


   @slow
   def     _SCREAMING_SNAKE_CASE    (     self			:					int  ):

     '''simple docstring'''
     for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       __A						      =			NatModel.from_pretrained(_lowerCamelCase  )
       self.assertIsNotNone(_lowerCamelCase  )






   def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[int]  ):

     '''simple docstring'''
     __A      ,   __A						      =			self.model_tester.prepare_config_and_inputs_for_common()

     __A						      =			_config_zero_init(_lowerCamelCase  )
     for model_class in self.all_model_classes:
       __A						      =			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  snake_case		(							unittest.TestCase						):


   '''simple docstring'''

   @cached_property
   def     _SCREAMING_SNAKE_CASE    (     self			:					str  ):

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






   @slow
   def     _SCREAMING_SNAKE_CASE    (     self			:					str  ):

     '''simple docstring'''
     __A						      =			NatForImageClassification.from_pretrained('''shi-labs/nat-mini-in1k-224'''  ).to(_lowerCamelCase  )
     __A						      =			self.default_image_processor

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

     # forward pass
     with torch.no_grad():
       __A						      =			model(**_lowerCamelCase  )

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





@require_torch
@require_natten
class  snake_case		(							unittest.TestCase       ,     _lowerCAmelCase						):


   '''simple docstring'''

   A_						:     Optional[Any]									=			(NatBackbone,) if is_torch_available() else ()
   A_						:     Any									=			NatConfig


   def     _SCREAMING_SNAKE_CASE    (     self			:					Optional[int]  ):

     '''simple docstring'''
     __A						      =			NatModelTester(self  )












"""simple docstring"""


from collections import OrderedDict
from dataclasses import fields
from typing import Any, Tuple

import numpy as np

from .import_utils import is_torch_available



def  lowerCAmelCase    (   __UpperCamelCase			):




  """simple docstring"""







  if is_torch_available():
    import torch

    if isinstance(__UpperCamelCase		,	torch.Tensor			):
      return True

  return isinstance(__UpperCamelCase		,	np.ndarray			)





class  snake_case		(							_lowerCAmelCase						):


   '''simple docstring'''

   def     _SCREAMING_SNAKE_CASE    (     self			:					Dict  ):

     '''simple docstring'''
     __A						      =			fields(self  )

     # Safety and consistency checks
     if not len(_lowerCamelCase  ):
       raise ValueError(f'{self.__class__.__name__} has no fields.'  )

     __A						      =			getattr(self,	class_fields[0].name  )
     __A						      =			all(getattr(self,	field.name  ) is None for field in class_fields[1:]  )

     if other_fields_are_none and isinstance(_lowerCamelCase,	_lowerCamelCase  ):
       for key, value in first_field.items():
         __A						      =			value
     else:
       for field in class_fields:
         __A						      =			getattr(self,	field.name  )
         if v is not None:
           __A						      =			v


   def __delitem__(     self			:					Dict,	*_lowerCamelCase			:					int,	**_lowerCamelCase			:					Any  ):

     '''simple docstring'''
     raise Exception(f'You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.'  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[Any],	*_lowerCamelCase			:					List[str],	**_lowerCamelCase			:					str  ):

     '''simple docstring'''
     raise Exception(f'You cannot use ``setdefault`` on a {self.__class__.__name__} instance.'  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					Any,	*_lowerCamelCase			:					Optional[Any],	**_lowerCamelCase			:					str  ):

     '''simple docstring'''
     raise Exception(f'You cannot use ``pop`` on a {self.__class__.__name__} instance.'  )


   def     _SCREAMING_SNAKE_CASE    (     self			:					List[str],	*_lowerCamelCase			:					Any,	**_lowerCamelCase			:					Optional[int]  ):

     '''simple docstring'''
     raise Exception(f'You cannot use ``update`` on a {self.__class__.__name__} instance.'  )


   def __getitem__(     self			:					int,	_lowerCamelCase			:					List[Any]  ):

     '''simple docstring'''
     if isinstance(_lowerCamelCase,	_lowerCamelCase  ):
       __A						      =			dict(self.items()  )
       return inner_dict[k]
     else:
       return self.to_tuple()[k]


   def __setattr__(     self			:					Optional[int],	_lowerCamelCase			:					Tuple,	_lowerCamelCase			:					Union[str, Any]  ):

     '''simple docstring'''
     if name in self.keys() and value is not None:
       # Don't call self.__setitem__ to avoid recursion errors
       super().__setitem__(_lowerCamelCase,	_lowerCamelCase  )
     super().__setattr__(_lowerCamelCase,	_lowerCamelCase  )


   def __setitem__(     self			:					List[str],	_lowerCamelCase			:					Optional[int],	_lowerCamelCase			:					str  ):

     '''simple docstring'''
     # Will raise a KeyException if needed
     super().__setitem__(_lowerCamelCase,	_lowerCamelCase  )
     # Don't call self.__setattr__ to avoid recursion errors
     super().__setattr__(_lowerCamelCase,	_lowerCamelCase  )






   def     _SCREAMING_SNAKE_CASE    (     self			:					Any  ):

     '''simple docstring'''
     return tuple(self[k] for k in self.keys()  )




