




'''simple docstring'''




import os
import tempfile
import unittest

import numpy as np
import torch

from diffusers.models import ModelMixin, UNetaDConditionModel
from diffusers.models.attention_processor import AttnProcessor, LoRAAttnProcessor
from diffusers.utils import (
    floats_tensor,
    logging,
    skip_mps,
    torch_device,
)
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

snake_case__          =						logging.get_logger(__name__)




def 						snake_case__    ( lowerCamelCase__				:   Optional[Any]			,		lowerCamelCase__				:   bool = True			)   ->   List[Any]:
					A_ :     List[Any] =							{}
					for name in model.attn_processors.keys():
										A_ :     Optional[int] =							name.endswith('''attn2.processor'''			) and not (
										    name.startswith('''transformer_in'''			) or '''temp_attentions''' in name.split('''.'''			)
										)
										A_ :     Optional[int] =							model.config.cross_attention_dim if has_cross_attention else None
										if name.startswith('''mid_block'''			):
															A_ :     List[Any] =							model.config.block_out_channels[-1]
										elif name.startswith('''up_blocks'''			):
															A_ :     str =							int(name[len('''up_blocks.'''			)]			)
															A_ :     List[str] =							list(reversed(model.config.block_out_channels			)			)[block_id]
										elif name.startswith('''down_blocks'''			):
															A_ :     Optional[int] =							int(name[len('''down_blocks.'''			)]			)
															A_ :     Union[str, Any] =							model.config.block_out_channels[block_id]
										elif name.startswith('''transformer_in'''			):
															# Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148
															A_ :     int =							8 * model.config.attention_head_dim

										A_ :     int =							LoRAAttnProcessor(hidden_size=lowerCamelCase__			,		cross_attention_dim=lowerCamelCase__			)
										A_ :     List[str] =							lora_attn_procs[name].to(model.device			)

										if mock_weights:
															# add 1 to weights to mock trained weights
															with torch.no_grad():
																				lora_attn_procs[name].to_q_lora.up.weight += 1
																				lora_attn_procs[name].to_k_lora.up.weight += 1
																				lora_attn_procs[name].to_v_lora.up.weight += 1
																				lora_attn_procs[name].to_out_lora.up.weight += 1

					return lora_attn_procs






@skip_mps
class     UpperCamelCase_     (a__,      a__,      unittest.TestCase     ):







						"""simple docstring"""


						_lowerCAmelCase    						=					UNetaDConditionModel
						_lowerCAmelCase    						=					'sample'







						@property
						def       _a    (     self      :       int    ):





											"""simple docstring"""







											A_ :     int =							4
											A_ :     Optional[Any] =							4
											A_ :     Any =							4
											A_ :     Dict =							(32, 32)

											A_ :     List[str] =							floats_tensor((batch_size, num_channels, num_frames) + sizes    ).to(_lowerCamelCase    )
											A_ :     List[Any] =							torch.tensor([10]    ).to(_lowerCamelCase    )
											A_ :     Tuple =							floats_tensor((batch_size, 4, 32)    ).to(_lowerCamelCase    )

											return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}







						@property
						def       _a    (     self      :       Tuple    ):





											"""simple docstring"""







											return (4, 4, 32, 32)







						@property
						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											return (4, 4, 32, 32)







						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											A_ :     str =							{
											    '''block_out_channels''': (32, 64),
											    '''down_block_types''': (
											        '''CrossAttnDownBlock3D''',
											        '''DownBlock3D''',
											    ),
											    '''up_block_types''': ('''UpBlock3D''', '''CrossAttnUpBlock3D'''),
											    '''cross_attention_dim''': 32,
											    '''attention_head_dim''': 8,
											    '''out_channels''': 4,
											    '''in_channels''': 4,
											    '''layers_per_block''': 1,
											    '''sample_size''': 32,
											}
											A_ :     int =							self.dummy_input
											return init_dict, inputs_dict







						@unittest.skipIf(
						    torch_device != '''cuda''' or not is_xformers_available()							, reason='''XFormers attention is only available with CUDA and `xformers` installed'''							, )
						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											A_ ,A_ :     Dict =							self.prepare_init_args_and_inputs_for_common()
											A_ :     int =							self.model_class(**_lowerCamelCase    )

											model.enable_xformers_memory_efficient_attention()

											assert (
											    model.mid_block.attentions[0].transformer_blocks[0].attna.processor.__class__.__name__
											    == "XFormersAttnProcessor"
											), "xformers is not enabled"







						def       _a    (     self      :       int    ):





											"""simple docstring"""







											A_ ,A_ :     Any =							self.prepare_init_args_and_inputs_for_common()

											A_ :     List[str] =							32

											A_ :     Dict =							self.model_class(**_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()

											with torch.no_grad():
																A_ :     List[Any] =							model(**_lowerCamelCase    )

																if isinstance(_lowerCamelCase							, _lowerCamelCase    ):
																					A_ :     int =							output.sample

											self.assertIsNotNone(_lowerCamelCase    )
											A_ :     Tuple =							inputs_dict['''sample'''].shape
											self.assertEqual(output.shape							, _lowerCamelCase							, '''Input and output shapes do not match'''    )







						def       _a    (     self      :       str    ):





											"""simple docstring"""







											A_ ,A_ :     List[Any] =							self.prepare_init_args_and_inputs_for_common()
											A_ :     Union[str, Any] =							self.model_class(**_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()

											with torch.no_grad():
																# Warmup pass when using mps (see #372)
																if torch_device == "mps" and isinstance(_lowerCamelCase							, _lowerCamelCase    ):
																					model(**self.dummy_input    )

																A_ :     List[Any] =							model(**_lowerCamelCase    )
																if isinstance(_lowerCamelCase							, _lowerCamelCase    ):
																					A_ :     Optional[int] =							first.sample

																A_ :     Optional[int] =							model(**_lowerCamelCase    )
																if isinstance(_lowerCamelCase							, _lowerCamelCase    ):
																					A_ :     Any =							second.sample

											A_ :     Any =							first.cpu().numpy()
											A_ :     Tuple =							second.cpu().numpy()
											A_ :     Tuple =							out_a[~np.isnan(_lowerCamelCase    )]
											A_ :     List[str] =							out_a[~np.isnan(_lowerCamelCase    )]
											A_ :     Dict =							np.amax(np.abs(out_a - out_a    )    )
											self.assertLessEqual(_lowerCamelCase							, 1E-5    )







						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											A_ ,A_ :     Union[str, Any] =							self.prepare_init_args_and_inputs_for_common()

											A_ :     List[Any] =							8

											A_ :     str =							self.model_class(**_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()

											model.set_attention_slice('''auto'''    )
											with torch.no_grad():
																A_ :     Optional[Any] =							model(**_lowerCamelCase    )
											assert output is not None

											model.set_attention_slice('''max'''    )
											with torch.no_grad():
																A_ :     List[Any] =							model(**_lowerCamelCase    )
											assert output is not None

											model.set_attention_slice(2    )
											with torch.no_grad():
																A_ :     Any =							model(**_lowerCamelCase    )
											assert output is not None







						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







											A_ ,A_ :     Optional[Any] =							self.prepare_init_args_and_inputs_for_common()

											A_ :     List[Any] =							8

											A_ :     Any =							self.model_class(**_lowerCamelCase    )
											model.to(_lowerCamelCase    )

											with torch.no_grad():
																A_ :     List[str] =							model(**_lowerCamelCase    ).sample

											A_ :     Tuple =							create_lora_layers(_lowerCamelCase    )

											# make sure we can set a list of attention processors
											model.set_attn_processor(_lowerCamelCase    )
											model.to(_lowerCamelCase    )

											# test that attn processors can be set to itself
											model.set_attn_processor(model.attn_processors    )

											with torch.no_grad():
																A_ :     str =							model(**_lowerCamelCase							, cross_attention_kwargs={'''scale''': 0.0}    ).sample
																A_ :     Union[str, Any] =							model(**_lowerCamelCase							, cross_attention_kwargs={'''scale''': 0.5}    ).sample
																A_ :     List[str] =							model(**_lowerCamelCase							, cross_attention_kwargs={'''scale''': 0.5}    ).sample

											assert (samplea - samplea).abs().max() < 3E-3
											assert (samplea - samplea).abs().max() < 3E-3

											# sample 2 and sample 3 should be different
											assert (samplea - samplea).abs().max() > 3E-3







						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											A_ ,A_ :     List[Any] =							self.prepare_init_args_and_inputs_for_common()

											A_ :     Tuple =							8

											torch.manual_seed(0    )
											A_ :     List[str] =							self.model_class(**_lowerCamelCase    )
											model.to(_lowerCamelCase    )

											with torch.no_grad():
																A_ :     Tuple =							model(**_lowerCamelCase    ).sample

											A_ :     List[Any] =							create_lora_layers(_lowerCamelCase    )
											model.set_attn_processor(_lowerCamelCase    )

											with torch.no_grad():
																A_ :     str =							model(**_lowerCamelCase							, cross_attention_kwargs={'''scale''': 0.5}    ).sample

											with tempfile.TemporaryDirectory() as tmpdirname:
																model.save_attn_procs(_lowerCamelCase    )
																self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase							, '''pytorch_lora_weights.bin'''    )    )    )
																torch.manual_seed(0    )
																A_ :     List[str] =							self.model_class(**_lowerCamelCase    )
																new_model.to(_lowerCamelCase    )
																new_model.load_attn_procs(_lowerCamelCase    )

											with torch.no_grad():
																A_ :     Union[str, Any] =							new_model(**_lowerCamelCase							, cross_attention_kwargs={'''scale''': 0.5}    ).sample

											assert (sample - new_sample).abs().max() < 1E-3

											# LoRA and no LoRA should NOT be the same
											assert (sample - old_sample).abs().max() > 1E-4







						def       _a    (     self      :       int    ):





											"""simple docstring"""







											A_ ,A_ :     Any =							self.prepare_init_args_and_inputs_for_common()

											A_ :     List[str] =							8

											torch.manual_seed(0    )
											A_ :     Tuple =							self.model_class(**_lowerCamelCase    )
											model.to(_lowerCamelCase    )

											with torch.no_grad():
																A_ :     int =							model(**_lowerCamelCase    ).sample

											A_ :     str =							create_lora_layers(_lowerCamelCase    )
											model.set_attn_processor(_lowerCamelCase    )

											with torch.no_grad():
																A_ :     str =							model(**_lowerCamelCase							, cross_attention_kwargs={'''scale''': 0.5}    ).sample

											with tempfile.TemporaryDirectory() as tmpdirname:
																model.save_attn_procs(_lowerCamelCase							, safe_serialization=_lowerCamelCase    )
																self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase							, '''pytorch_lora_weights.safetensors'''    )    )    )
																torch.manual_seed(0    )
																A_ :     str =							self.model_class(**_lowerCamelCase    )
																new_model.to(_lowerCamelCase    )
																new_model.load_attn_procs(_lowerCamelCase    )

											with torch.no_grad():
																A_ :     List[Any] =							new_model(**_lowerCamelCase							, cross_attention_kwargs={'''scale''': 0.5}    ).sample

											assert (sample - new_sample).abs().max() < 3E-3

											# LoRA and no LoRA should NOT be the same
											assert (sample - old_sample).abs().max() > 1E-4







						def       _a    (     self      :       List[Any]    ):





											"""simple docstring"""







											A_ ,A_ :     Tuple =							self.prepare_init_args_and_inputs_for_common()

											A_ :     Any =							8

											torch.manual_seed(0    )
											A_ :     List[Any] =							self.model_class(**_lowerCamelCase    )
											model.to(_lowerCamelCase    )

											A_ :     Union[str, Any] =							create_lora_layers(_lowerCamelCase							, mock_weights=_lowerCamelCase    )
											model.set_attn_processor(_lowerCamelCase    )
											# Saving as torch, properly reloads with directly filename
											with tempfile.TemporaryDirectory() as tmpdirname:
																model.save_attn_procs(_lowerCamelCase    )
																self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase							, '''pytorch_lora_weights.bin'''    )    )    )
																torch.manual_seed(0    )
																A_ :     Optional[int] =							self.model_class(**_lowerCamelCase    )
																new_model.to(_lowerCamelCase    )
																new_model.load_attn_procs(_lowerCamelCase							, weight_name='''pytorch_lora_weights.bin'''    )







						def       _a    (     self      :       List[Any]    ):





											"""simple docstring"""







											A_ ,A_ :     Any =							self.prepare_init_args_and_inputs_for_common()

											A_ :     Union[str, Any] =							8

											torch.manual_seed(0    )
											A_ :     List[Any] =							self.model_class(**_lowerCamelCase    )
											model.to(_lowerCamelCase    )

											A_ :     str =							create_lora_layers(_lowerCamelCase							, mock_weights=_lowerCamelCase    )
											model.set_attn_processor(_lowerCamelCase    )
											# Saving as torch, properly reloads with directly filename
											with tempfile.TemporaryDirectory() as tmpdirname:
																model.save_attn_procs(_lowerCamelCase    )
																self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase							, '''pytorch_lora_weights.bin'''    )    )    )
																torch.manual_seed(0    )
																A_ :     List[Any] =							self.model_class(**_lowerCamelCase    )
																new_model.to(_lowerCamelCase    )
																with self.assertRaises(_lowerCamelCase    ) as e:
																					new_model.load_attn_procs(_lowerCamelCase							, use_safetensors=_lowerCamelCase    )
																self.assertIn('''Error no file named pytorch_lora_weights.safetensors'''							, str(e.exception    )    )







						def       _a    (     self      :       Tuple    ):





											"""simple docstring"""







											A_ ,A_ :     List[Any] =							self.prepare_init_args_and_inputs_for_common()

											A_ :     Optional[Any] =							8

											torch.manual_seed(0    )
											A_ :     List[str] =							self.model_class(**_lowerCamelCase    )
											model.to(_lowerCamelCase    )

											with torch.no_grad():
																A_ :     str =							model(**_lowerCamelCase    ).sample

											A_ :     Any =							create_lora_layers(_lowerCamelCase    )
											model.set_attn_processor(_lowerCamelCase    )

											with torch.no_grad():
																A_ :     Union[str, Any] =							model(**_lowerCamelCase							, cross_attention_kwargs={'''scale''': 0.0}    ).sample

											model.set_attn_processor(AttnProcessor()    )

											with torch.no_grad():
																A_ :     Any =							model(**_lowerCamelCase    ).sample

											assert (sample - new_sample).abs().max() < 1E-4
											assert (sample - old_sample).abs().max() < 3E-3







						@unittest.skipIf(
						    torch_device != '''cuda''' or not is_xformers_available()							, reason='''XFormers attention is only available with CUDA and `xformers` installed'''							, )
						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







											A_ ,A_ :     Tuple =							self.prepare_init_args_and_inputs_for_common()

											A_ :     Optional[int] =							4

											torch.manual_seed(0    )
											A_ :     Dict =							self.model_class(**_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											A_ :     str =							create_lora_layers(_lowerCamelCase    )
											model.set_attn_processor(_lowerCamelCase    )

											# default
											with torch.no_grad():
																A_ :     Optional[int] =							model(**_lowerCamelCase    ).sample

																model.enable_xformers_memory_efficient_attention()
																A_ :     Optional[int] =							model(**_lowerCamelCase    ).sample

																model.disable_xformers_memory_efficient_attention()
																A_ :     Dict =							model(**_lowerCamelCase    ).sample

											assert (sample - on_sample).abs().max() < 1E-4
											assert (sample - off_sample).abs().max() < 1E-4







						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







											A_ ,A_ :     Optional[Any] =							self.prepare_init_args_and_inputs_for_common()
											A_ :     List[Any] =							32

											A_ :     str =							self.model_class(**_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()

											with torch.no_grad():
																A_ :     str =							model(**_lowerCamelCase    )[0]

											model.enable_forward_chunking()
											with torch.no_grad():
																A_ :     Any =							model(**_lowerCamelCase    )[0]

											self.assertEqual(output.shape							, output_a.shape							, '''Shape doesn\'t match'''    )
											assert np.abs(output.cpu() - output_a.cpu()    ).max() < 1E-2


# (todo: sayakpaul) implement SLOW tests.









'''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     UpperCamelCase_     :







						"""simple docstring"""


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





											"""simple docstring"""







											A_ :     Any =							parent
											A_ :     Union[str, Any] =							batch_size
											A_ :     Optional[int] =							image_size
											A_ :     int =							patch_size
											A_ :     str =							num_channels
											A_ :     Dict =							embed_dim
											A_ :     int =							depths
											A_ :     str =							num_heads
											A_ :     str =							kernel_size
											A_ :     Tuple =							mlp_ratio
											A_ :     Tuple =							qkv_bias
											A_ :     Any =							hidden_dropout_prob
											A_ :     List[Any] =							attention_probs_dropout_prob
											A_ :     List[str] =							drop_path_rate
											A_ :     Optional[Any] =							hidden_act
											A_ :     Any =							patch_norm
											A_ :     str =							layer_norm_eps
											A_ :     int =							initializer_range
											A_ :     Optional[int] =							is_training
											A_ :     Optional[int] =							scope
											A_ :     List[Any] =							use_labels
											A_ :     int =							num_labels
											A_ :     Union[str, Any] =							out_features
											A_ :     Optional[Any] =							out_indices







						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											A_ :     Dict =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

											A_ :     int =							None
											if self.use_labels:
																A_ :     Union[str, Any] =							ids_tensor([self.batch_size]							, self.num_labels    )

											A_ :     int =							self.get_config()

											return config, pixel_values, labels







						def       _a    (     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       _a    (     self      :       int							, _lowerCamelCase      :       str							, _lowerCamelCase      :       Dict							, _lowerCamelCase      :       Dict    ):





											"""simple docstring"""







											A_ :     str =							NatModel(config=_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()
											A_ :     Union[str, Any] =							model(_lowerCamelCase    )

											A_ :     List[str] =							(config.image_size // config.patch_size) // (2 ** (len(config.depths    ) - 1))
											A_ :     Optional[int] =							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       _a    (     self      :       str							, _lowerCamelCase      :       Tuple							, _lowerCamelCase      :       int							, _lowerCamelCase      :       Union[str, Any]    ):





											"""simple docstring"""







											A_ :     int =							NatForImageClassification(_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()
											A_ :     str =							model(_lowerCamelCase							, labels=_lowerCamelCase    )
											self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.num_labels)    )

											# test greyscale images
											A_ :     str =							1
											A_ :     Dict =							NatForImageClassification(_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()

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







						def       _a    (     self      :       List[Any]							, _lowerCamelCase      :       List[str]							, _lowerCamelCase      :       Any							, _lowerCamelCase      :       Tuple    ):





											"""simple docstring"""







											A_ :     Union[str, Any] =							NatBackbone(config=_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()
											A_ :     Optional[int] =							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_ :     Tuple =							None
											A_ :     List[Any] =							NatBackbone(config=_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()
											A_ :     Optional[int] =							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       _a    (     self      :       List[Any]    ):





											"""simple docstring"""







											A_ :     Optional[int] =							self.prepare_config_and_inputs()
											A_ ,A_ ,A_ :     Any =							config_and_inputs
											A_ :     Union[str, Any] =							{'''pixel_values''': pixel_values}
											return config, inputs_dict






@require_natten
@require_torch
class     UpperCamelCase_     (a__,      a__,      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       _a    (     self      :       Any    ):





											"""simple docstring"""







											A_ :     Tuple =							NatModelTester(self    )
											A_ :     Any =							ConfigTester(self							, config_class=_lowerCamelCase							, embed_dim=37    )







						def       _a    (     self      :       Union[str, 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       _a    (     self      :       str    ):





											"""simple docstring"""







											return







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





											"""simple docstring"""







											A_ :     Union[str, Any] =							self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_lowerCamelCase    )







						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											A_ :     Dict =							self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*_lowerCamelCase    )







						def       _a    (     self      :       str    ):





											"""simple docstring"""







											A_ :     str =							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       _a    (     self      :       str    ):





											"""simple docstring"""







											pass







						@unittest.skip(reason='''Nat does not use feedforward chunking'''    )
						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											pass







						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											A_ ,A_ :     List[Any] =							self.model_tester.prepare_config_and_inputs_for_common()

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







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





											"""simple docstring"""







											A_ ,A_ :     Tuple =							self.model_tester.prepare_config_and_inputs_for_common()

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

																A_ :     List[Any] =							['''pixel_values''']
																self.assertListEqual(arg_names[:1]							, _lowerCamelCase    )







						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







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







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





											"""simple docstring"""







											A_ :     Tuple =							model_class(_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()

											with torch.no_grad():
																A_ :     List[str] =							model(**self._prepare_for_class(_lowerCamelCase							, _lowerCamelCase    )    )

											A_ :     List[Any] =							outputs.hidden_states

											A_ :     List[str] =							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_ :     int =							(
											    config.patch_size
											    if isinstance(config.patch_size							, collections.abc.Iterable    )
											    else (config.patch_size, config.patch_size)
											)

											A_ :     Optional[Any] =							image_size[0] // patch_size[0]
											A_ :     Optional[int] =							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_ :     Tuple =							outputs.reshaped_hidden_states
																self.assertEqual(len(_lowerCamelCase    )							, _lowerCamelCase    )

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





											"""simple docstring"""







											A_ ,A_ :     List[str] =							self.model_tester.prepare_config_and_inputs_for_common()

											A_ :     List[str] =							(
											    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_ :     Tuple =							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_ :     Dict =							True

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







						@slow
						def       _a    (     self      :       Tuple    ):





											"""simple docstring"""







											for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																A_ :     int =							NatModel.from_pretrained(_lowerCamelCase    )
																self.assertIsNotNone(_lowerCamelCase    )







						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







											A_ ,A_ :     List[str] =							self.model_tester.prepare_config_and_inputs_for_common()

											A_ :     Tuple =							_config_zero_init(_lowerCamelCase    )
											for model_class in self.all_model_classes:
																A_ :     Union[str, Any] =							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     UpperCamelCase_     (unittest.TestCase     ):







						"""simple docstring"""


						@cached_property
						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







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







						@slow
						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											A_ :     Optional[Any] =							NatForImageClassification.from_pretrained('''shi-labs/nat-mini-in1k-224'''    ).to(_lowerCamelCase    )
											A_ :     List[Any] =							self.default_image_processor

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

											# forward pass
											with torch.no_grad():
																A_ :     Union[str, Any] =							model(**_lowerCamelCase    )

											# verify the logits
											A_ :     int =							torch.Size((1, 1000)    )
											self.assertEqual(outputs.logits.shape							, _lowerCamelCase    )
											A_ :     int =							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     UpperCamelCase_     (unittest.TestCase,      a__     ):







						"""simple docstring"""


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







						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											A_ :     Union[str, Any] =							NatModelTester(self    )



