



import collections
import inspect
import unittest

from transformers import DinatConfig
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 DinatBackbone, DinatForImageClassification, DinatModel
		from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor


class 	a__       :


					"""simple docstring"""



					def __init__(		self     :str      ,  lowercase__     :Optional[Any]      ,  lowercase__     :Tuple=13      ,  lowercase__     :str=64      ,  lowercase__     :int=4      ,  lowercase__     :Tuple=3      ,  lowercase__     :List[str]=16      ,  lowercase__     :Any=[1, 2, 1]      ,  lowercase__     :Tuple=[2, 4, 8]      ,  lowercase__     :int=3      ,  lowercase__     :Any=[[3], [1, 2], [1]]      ,  lowercase__     :Tuple=2.0      ,  lowercase__     :Optional[Any]=True      ,  lowercase__     :Optional[Any]=0.0      ,  lowercase__     :str=0.0      ,  lowercase__     :str=0.1      ,  lowercase__     :str="gelu"      ,  lowercase__     :List[str]=True      ,  lowercase__     :List[Any]=0.02      ,  lowercase__     :Optional[Any]=1E-5      ,  lowercase__     :str=True      ,  lowercase__     :List[str]=None      ,  lowercase__     :str=True      ,  lowercase__     :Optional[int]=10      ,  lowercase__     :Any=["stage1", "stage2"]      ,  lowercase__     :int=[1, 2]      ,  ):
										lowercase						   =       parent
										lowercase						   =       batch_size
										lowercase						   =       image_size
										lowercase						   =       patch_size
										lowercase						   =       num_channels
										lowercase						   =       embed_dim
										lowercase						   =       depths
										lowercase						   =       num_heads
										lowercase						   =       kernel_size
										lowercase						   =       dilations
										lowercase						   =       mlp_ratio
										lowercase						   =       qkv_bias
										lowercase						   =       hidden_dropout_prob
										lowercase						   =       attention_probs_dropout_prob
										lowercase						   =       drop_path_rate
										lowercase						   =       hidden_act
										lowercase						   =       patch_norm
										lowercase						   =       layer_norm_eps
										lowercase						   =       initializer_range
										lowercase						   =       is_training
										lowercase						   =       scope
										lowercase						   =       use_labels
										lowercase						   =       num_labels
										lowercase						   =       out_features
										lowercase						   =       out_indices

					def 				__UpperCAmelCase		(		self     :str    ):
										lowercase						   =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

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

										lowercase						   =       self.get_config()

										return config, pixel_values, labels

					def 				__UpperCAmelCase		(		self     :str    ):
										return DinatConfig(
										    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      ,  dilations=self.dilations      ,  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     :Optional[Any]      ,  lowercase__     :Tuple      ,  lowercase__     :Optional[Any]      ,  lowercase__     :Optional[Any]    ):
										lowercase						   =       DinatModel(config=_lowerCamelCase    )
										model.to(_lowerCamelCase    )
										model.eval()
										lowercase						   =       model(_lowerCamelCase    )

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

										# test greyscale images
										lowercase						   =       1
										lowercase						   =       DinatForImageClassification(_lowerCamelCase    )
										model.to(_lowerCamelCase    )
										model.eval()

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

					def 				__UpperCAmelCase		(		self     :Tuple      ,  lowercase__     :Tuple      ,  lowercase__     :Tuple      ,  lowercase__     :Tuple    ):
										lowercase						   =       DinatBackbone(config=_lowerCamelCase    )
										model.to(_lowerCamelCase    )
										model.eval()
										lowercase						   =       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
										lowercase						   =       None
										lowercase						   =       DinatBackbone(config=_lowerCamelCase    )
										model.to(_lowerCamelCase    )
										model.eval()
										lowercase						   =       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 				__UpperCAmelCase		(		self     :str    ):
										lowercase						   =       self.prepare_config_and_inputs()
										lowercase						   =       config_and_inputs
										lowercase						   =       {'''pixel_values''': pixel_values}
										return config, inputs_dict




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


					"""simple docstring"""



					A__  :				List[str]					 =				(
					    (
					        DinatModel,
					        DinatForImageClassification,
					        DinatBackbone,
					    )
					    if is_torch_available()
					    else ()
					)
					A__  :				List[Any]					 =				(
					    {'''feature-extraction''': DinatModel, '''image-classification''': DinatForImageClassification}
					    if is_torch_available()
					    else {}
					)
					A__  :				Optional[int]					 =				False

					A__  :				str					 =				False
					A__  :				Dict					 =				False
					A__  :				Optional[int]					 =				False
					A__  :				List[str]					 =				False

					def 				__UpperCAmelCase		(		self     :Optional[Any]    ):
										lowercase						   =       DinatModelTester(self    )
										lowercase						   =       ConfigTester(self      ,  config_class=_lowerCamelCase      ,  embed_dim=37    )

					def 				__UpperCAmelCase		(		self     :Any    ):
										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     :Optional[Any]    ):
										return

					def 				__UpperCAmelCase		(		self     :List[Any]    ):
										lowercase						   =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_lowerCamelCase    )

					def 				__UpperCAmelCase		(		self     :List[str]    ):
										lowercase						   =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*_lowerCamelCase    )

					def 				__UpperCAmelCase		(		self     :Optional[Any]    ):
										lowercase						   =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_backbone(*_lowerCamelCase    )

					@unittest.skip(reason='Dinat does not use inputs_embeds'    )
					def 				__UpperCAmelCase		(		self     :Tuple    ):
										pass

					@unittest.skip(reason='Dinat does not use feedforward chunking'    )
					def 				__UpperCAmelCase		(		self     :List[Any]    ):
										pass

					def 				__UpperCAmelCase		(		self     :Optional[int]    ):
										lowercase						   =       self.model_tester.prepare_config_and_inputs_for_common()

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

					def 				__UpperCAmelCase		(		self     :Union[str, Any]    ):
										lowercase						   =       self.model_tester.prepare_config_and_inputs_for_common()

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

															lowercase						   =       ['''pixel_values''']
															self.assertListEqual(arg_names[:1]      ,  _lowerCamelCase    )

					def 				__UpperCAmelCase		(		self     :Optional[int]    ):
										self.skipTest('Dinat\'s attention operation is handled entirely by NATTEN.'    )

					def 				__UpperCAmelCase		(		self     :Optional[Any]      ,  lowercase__     :Optional[Any]      ,  lowercase__     :List[Any]      ,  lowercase__     :List[Any]      ,  lowercase__     :str    ):
										lowercase						   =       model_class(_lowerCamelCase    )
										model.to(_lowerCamelCase    )
										model.eval()

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

										lowercase						   =       outputs.hidden_states

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

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

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

															lowercase						   =       reshaped_hidden_states[0].shape
															lowercase						   =       (
															    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 				__UpperCAmelCase		(		self     :int    ):
										lowercase						   =       self.model_tester.prepare_config_and_inputs_for_common()

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

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

					@slow
					def 				__UpperCAmelCase		(		self     :List[Any]    ):







										for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															lowercase						   =       DinatModel.from_pretrained(_lowerCamelCase    )
															self.assertIsNotNone(_lowerCamelCase    )



					def 				__UpperCAmelCase		(		self     :str    ):
										lowercase						   =       self.model_tester.prepare_config_and_inputs_for_common()

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


					"""simple docstring"""



					@cached_property
					def 				__UpperCAmelCase		(		self     :List[str]    ):
										return AutoImageProcessor.from_pretrained('shi-labs/dinat-mini-in1k-224'    ) if is_vision_available() else None



					@slow
					def 				__UpperCAmelCase		(		self     :List[str]    ):
										lowercase						   =       DinatForImageClassification.from_pretrained('shi-labs/dinat-mini-in1k-224'    ).to(_lowerCamelCase    )
										lowercase						   =       self.default_image_processor

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

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

										# verify the logits
										lowercase						   =       torch.Size((1, 1000)    )
										self.assertEqual(outputs.logits.shape      ,  _lowerCamelCase    )
										lowercase						   =       torch.tensor([-0.1545, -0.7667, 0.4642]    ).to(_lowerCamelCase    )
										self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,  _lowerCamelCase      ,  atol=1E-4    )    )




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


					"""simple docstring"""



					A__  :				List[str]					 =				(DinatBackbone,) if is_torch_available() else ()
					A__  :				List[Any]					 =				DinatConfig

					def 				__UpperCAmelCase		(		self     :Optional[Any]    ):
										lowercase						   =       DinatModelTester(self    )












'''simple docstring'''




from ...configuration_utils import PretrainedConfig
from ...utils import logging


snake_case__ 							=    logging.get_logger(__name__)

snake_case__ 							=    {
    """uw-madison/nystromformer-512""": """https://huggingface.co/uw-madison/nystromformer-512/resolve/main/config.json""",
    # See all Nystromformer models at https://huggingface.co/models?filter=nystromformer
}






class      UpperCamelCase_  (a__  ):







						"""simple docstring"""


						_lowerCAmelCase         =       'nystromformer'







						def __init__(  self   :							Optional[Any]				,			_lowerCamelCase   :							str=30000				,			_lowerCamelCase   :							Optional[Any]=768				,			_lowerCamelCase   :							Optional[Any]=12				,			_lowerCamelCase   :							str=12				,			_lowerCamelCase   :							List[str]=3072				,			_lowerCamelCase   :							Any="gelu_new"				,			_lowerCamelCase   :							List[Any]=0.1				,			_lowerCamelCase   :							List[str]=0.1				,			_lowerCamelCase   :							List[str]=510				,			_lowerCamelCase   :							Optional[Any]=2				,			_lowerCamelCase   :							List[Any]=64				,			_lowerCamelCase   :							str=64				,			_lowerCamelCase   :							str=65				,			_lowerCamelCase   :							List[Any]=False				,			_lowerCamelCase   :							List[Any]=0.02				,			_lowerCamelCase   :							List[Any]=1E-5				,			_lowerCamelCase   :							Dict=1				,			_lowerCamelCase   :							Tuple=0				,			_lowerCamelCase   :							List[str]=2				,			**_lowerCamelCase   :							List[str]				,			):





											"""simple docstring"""







											A_			:  List[Any]			=				vocab_size
											A_			:  Tuple			=				max_position_embeddings
											A_			:  Dict			=				hidden_size
											A_			:  Optional[Any]			=				num_hidden_layers
											A_			:  Dict			=				num_attention_heads
											A_			:  Tuple			=				intermediate_size
											A_			:  Dict			=				hidden_act
											A_			:  List[Any]			=				hidden_dropout_prob
											A_			:  str			=				attention_probs_dropout_prob
											A_			:  Optional[Any]			=				initializer_range
											A_			:  List[str]			=				type_vocab_size
											A_			:  List[Any]			=				segment_means_seq_len
											A_			:  Dict			=				num_landmarks
											A_			:  int			=				conv_kernel_size
											A_			:  Optional[Any]			=				inv_coeff_init_option
											A_			:  str			=				layer_norm_eps
											super().__init__(pad_token_id=_lowerCamelCase				,			bos_token_id=_lowerCamelCase				,			eos_token_id=_lowerCamelCase				,			**_lowerCamelCase     )



