






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					:	Tuple  ,  lowerCAmelCase__					:	Union[str, Any]  ,  lowerCAmelCase__					:	Dict=1_3  ,  lowerCAmelCase__					:	Union[str, Any]=6_4  ,  lowerCAmelCase__					:	List[Any]=4  ,  lowerCAmelCase__					:	Union[str, Any]=3  ,  lowerCAmelCase__					:	Tuple=1_6  ,  lowerCAmelCase__					:	List[str]=[1, 2, 1]  ,  lowerCAmelCase__					:	Union[str, Any]=[2, 4, 8]  ,  lowerCAmelCase__					:	Tuple=3  ,  lowerCAmelCase__					:	Any=2.0  ,  lowerCAmelCase__					:	int=True  ,  lowerCAmelCase__					:	int=0.0  ,  lowerCAmelCase__					:	int=0.0  ,  lowerCAmelCase__					:	int=0.1  ,  lowerCAmelCase__					:	List[str]="gelu"  ,  lowerCAmelCase__					:	Optional[Any]=True  ,  lowerCAmelCase__					:	Any=0.02  ,  lowerCAmelCase__					:	List[str]=1e-5  ,  lowerCAmelCase__					:	List[Any]=True  ,  lowerCAmelCase__					:	Optional[int]=None  ,  lowerCAmelCase__					:	Optional[Any]=True  ,  lowerCAmelCase__					:	Dict=1_0  ,  lowerCAmelCase__					:	List[str]=["stage1", "stage2"]  ,  lowerCAmelCase__					:	Union[str, Any]=[1, 2]  ,  )							->	Optional[Any]:

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





     def     __UpperCamelCase						(		self					:	Optional[int])							->	Optional[Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

      SCREAMING_SNAKE_CASE				    =		self.get_config()

      return config, pixel_values, labels





     def     __UpperCamelCase						(		self					:	List[Any])							->	str:

      """simple docstring"""
      return NatConfig(
          num_labels=self.num_labels  ,  image_size=self.image_size  ,  patch_size=self.patch_size  ,  num_channels=self.num_channels  ,  embed_dim=self.embed_dim  ,  depths=self.depths  ,  num_heads=self.num_heads  ,  kernel_size=self.kernel_size  ,  mlp_ratio=self.mlp_ratio  ,  qkv_bias=self.qkv_bias  ,  hidden_dropout_prob=self.hidden_dropout_prob  ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,  drop_path_rate=self.drop_path_rate  ,  hidden_act=self.hidden_act  ,  patch_norm=self.patch_norm  ,  layer_norm_eps=self.layer_norm_eps  ,  initializer_range=self.initializer_range  ,  out_features=self.out_features  ,  out_indices=self.out_indices  ,  )





     def     __UpperCamelCase						(		self					:	Tuple  ,  lowerCAmelCase__					:	Dict  ,  lowerCAmelCase__					:	Dict  ,  lowerCAmelCase__					:	Dict)							->	str:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		NatModel(config=_A)
      model.to(_A)
      model.eval()
      SCREAMING_SNAKE_CASE				    =		model(_A)

      SCREAMING_SNAKE_CASE				    =		SCREAMING_SNAKE_CASE				    =		(config.image_size // config.patch_size) // (2 ** (len(config.depths) - 1))
      SCREAMING_SNAKE_CASE				    =		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[int]  ,  lowerCAmelCase__					:	str  ,  lowerCAmelCase__					:	Optional[int]  ,  lowerCAmelCase__					:	Optional[int])							->	Union[str, Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		NatForImageClassification(_A)
      model.to(_A)
      model.eval()
      SCREAMING_SNAKE_CASE				    =		model(_A  ,  labels=_A)
      self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels))

      # test greyscale images
      SCREAMING_SNAKE_CASE				    =		1
      SCREAMING_SNAKE_CASE				    =		NatForImageClassification(_A)
      model.to(_A)
      model.eval()

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





     def     __UpperCamelCase						(		self					:	Optional[Any]  ,  lowerCAmelCase__					:	int  ,  lowerCAmelCase__					:	str  ,  lowerCAmelCase__					:	List[Any])							->	int:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		NatBackbone(config=_A)
      model.to(_A)
      model.eval()
      SCREAMING_SNAKE_CASE				    =		model(_A)

      # 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
      SCREAMING_SNAKE_CASE				    =		None
      SCREAMING_SNAKE_CASE				    =		NatBackbone(config=_A)
      model.to(_A)
      model.eval()
      SCREAMING_SNAKE_CASE				    =		model(_A)

      # 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)							->	Tuple:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		self.prepare_config_and_inputs()
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		config_and_inputs
      SCREAMING_SNAKE_CASE				    =		{'pixel_values': pixel_values}
      return config, inputs_dict

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



     """simple docstring"""



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

     __snake_case		:				int  					=False
     __snake_case		:				str  					=False
     __snake_case		:				List[str]  					=False
     __snake_case		:				List[str]  					=False





     def     __UpperCamelCase						(		self					:	Optional[Any])							->	Tuple:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		NatModelTester(self)
      SCREAMING_SNAKE_CASE				    =		ConfigTester(self  ,  config_class=_A  ,  embed_dim=3_7)





     def     __UpperCamelCase						(		self					:	Any)							->	Any:

      """simple docstring"""
      self.create_and_test_config_common_properties()
      self.config_tester.create_and_test_config_to_json_string()
      self.config_tester.create_and_test_config_to_json_file()
      self.config_tester.create_and_test_config_from_and_save_pretrained()
      self.config_tester.create_and_test_config_with_num_labels()
      self.config_tester.check_config_can_be_init_without_params()
      self.config_tester.check_config_arguments_init()





     def     __UpperCamelCase						(		self					:	Dict)							->	Optional[int]:

      """simple docstring"""
      return





     def     __UpperCamelCase						(		self					:	str)							->	Union[str, Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*_A)





     def     __UpperCamelCase						(		self					:	str)							->	List[Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*_A)





     def     __UpperCamelCase						(		self					:	Union[str, Any])							->	List[Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_backbone(*_A)





     @unittest.skip(reason='Nat does not use inputs_embeds')
     def     __UpperCamelCase						(		self					:	int)							->	Union[str, Any]:

      """simple docstring"""
      pass





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

      """simple docstring"""
      pass





     def     __UpperCamelCase						(		self					:	Optional[Any])							->	int:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.model_tester.prepare_config_and_inputs_for_common()

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





     def     __UpperCamelCase						(		self					:	List[Any])							->	Optional[Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.model_tester.prepare_config_and_inputs_for_common()

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

       SCREAMING_SNAKE_CASE				    =		['pixel_values']
       self.assertListEqual(arg_names[:1]  ,  _A)





     def     __UpperCamelCase						(		self					:	Optional[int])							->	str:

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





     def     __UpperCamelCase						(		self					:	Tuple  ,  lowerCAmelCase__					:	Tuple  ,  lowerCAmelCase__					:	Dict  ,  lowerCAmelCase__					:	Any  ,  lowerCAmelCase__					:	List[Any])							->	Dict:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		model_class(_A)
      model.to(_A)
      model.eval()

      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**self._prepare_for_class(_A  ,  _A))

      SCREAMING_SNAKE_CASE				    =		outputs.hidden_states

      SCREAMING_SNAKE_CASE				    =		getattr(
          self.model_tester  ,  'expected_num_hidden_layers'  ,  len(self.model_tester.depths) + 1)
      self.assertEqual(len(_A)  ,  _A)

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

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

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





     def     __UpperCamelCase						(		self					:	Tuple)							->	Dict:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.model_tester.prepare_config_and_inputs_for_common()

      SCREAMING_SNAKE_CASE				    =		(
          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:
       SCREAMING_SNAKE_CASE				    =		True
       self.check_hidden_states_output(_A  ,  _A  ,  _A  ,  _A)

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

       self.check_hidden_states_output(_A  ,  _A  ,  _A  ,  _A)





     @slow
     def     __UpperCamelCase						(		self					:	Tuple)							->	List[str]:

      """simple docstring"""



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





     def     __UpperCamelCase						(		self					:	int)							->	Optional[Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self.model_tester.prepare_config_and_inputs_for_common()

      SCREAMING_SNAKE_CASE				    =		_config_zero_init(_A)
      for model_class in self.all_model_classes:
       SCREAMING_SNAKE_CASE				    =		model_class(config=_A)
       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					:	Optional[Any])							->	Optional[int]:

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





     @slow
     def     __UpperCamelCase						(		self					:	str)							->	Tuple:

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

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

      # forward pass
      with torch.no_grad():
       SCREAMING_SNAKE_CASE				    =		model(**_A)

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

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



     """simple docstring"""



     __snake_case		:				Any  					=(NatBackbone,) if is_torch_available() else ()
     __snake_case		:				Dict  					=NatConfig





     def     __UpperCamelCase						(		self					:	Tuple)							->	Optional[Any]:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		NatModelTester(self)






'''simple docstring'''





from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
 if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.27.0")):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from ...utils.dummy_torch_and_transformers_objects import (
     AudioLDMPipeline,
 )
else:
 from .pipeline_audioldm import AudioLDMPipeline






