


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      UpperCAmelCase_ :

   '''simple docstring'''




   def __init__(		self						,    __A						,    __A=13						,    __A=64						,    __A=4						,    __A=3						,    __A=16						,    __A=[1, 2, 1]						,    __A=[2, 4, 8]						,    __A=3						,    __A=[[3], [1, 2], [1]]						,    __A=2.0						,    __A=True						,    __A=0.0						,    __A=0.0						,    __A=0.1						,    __A="gelu"						,    __A=True						,    __A=0.02						,    __A=1e-5						,    __A=True						,    __A=None						,    __A=True						,    __A=10						,    __A=["stage1", "stage2"]						,    __A=[1, 2]						,    ):






     """simple docstring"""





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






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Any   =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

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

     lowerCamelCase			: int   =    self.get_config()

     return config, pixel_values, labels






   def _snake_case   (		self			):






     """simple docstring"""





     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 _snake_case   (		self						,    __A						,    __A						,    __A			):






     """simple docstring"""





     lowerCamelCase			: List[str]   =    DinatModel(config=__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: Optional[int]   =    model(__A			)

     lowerCamelCase			: List[str]   =    (config.image_size // config.patch_size) // (2 ** (len(config.depths			) - 1))
     lowerCamelCase			: List[Any]   =    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 _snake_case   (		self						,    __A						,    __A						,    __A			):






     """simple docstring"""





     lowerCamelCase			: List[Any]   =    DinatForImageClassification(__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: str   =    model(__A						,    labels=__A			)
     self.parent.assertEqual(result.logits.shape						,    (self.batch_size, self.num_labels)			)

     # test greyscale images
     lowerCamelCase			: Any   =    1
     lowerCamelCase			: Union[str, Any]   =    DinatForImageClassification(__A			)
     model.to(__A			)
     model.eval()

     lowerCamelCase			: Union[str, Any]   =    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
     lowerCamelCase			: int   =    model(__A			)
     self.parent.assertEqual(result.logits.shape						,    (self.batch_size, self.num_labels)			)






   def _snake_case   (		self						,    __A						,    __A						,    __A			):






     """simple docstring"""





     lowerCamelCase			: Union[str, Any]   =    DinatBackbone(config=__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: Tuple   =    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], 16, 16]			)

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

     # verify backbone works with out_features=None
     lowerCamelCase			: Union[str, Any]   =    None
     lowerCamelCase			: Optional[int]   =    DinatBackbone(config=__A			)
     model.to(__A			)
     model.eval()
     lowerCamelCase			: Any   =    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 _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: str   =    self.prepare_config_and_inputs()
     lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase			: List[Any]   =    config_and_inputs
     lowerCamelCase			: Optional[Any]   =    {"pixel_values": pixel_values}
     return config, inputs_dict



@require_natten
@require_torch
class      UpperCAmelCase_ ( UpperCamelCase		,			UpperCamelCase		,			unittest.TestCase   ):

   '''simple docstring'''




   __A			:				Any							=		(
       (
           DinatModel,
           DinatForImageClassification,
           DinatBackbone,
       )
       if is_torch_available()
       else ()
   )
   __A			:				Optional[int]							=		(
       {"feature-extraction": DinatModel, "image-classification": DinatForImageClassification}
       if is_torch_available()
       else {}
   )
   __A			:				Tuple							=		False

   __A			:				Tuple							=		False
   __A			:				Optional[Any]							=		False
   __A			:				Dict							=		False
   __A			:				str							=		False






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Optional[int]   =    DinatModelTester(self			)
     lowerCamelCase			: Union[str, Any]   =    ConfigTester(self						,    config_class=__A						,    embed_dim=37			)






   def _snake_case   (		self			):






     """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 _snake_case   (		self			):






     """simple docstring"""





     return






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Tuple   =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Tuple   =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*__A			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Optional[Any]   =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_backbone(*__A			)






   @unittest.skip(reason="Dinat does not use inputs_embeds"			)
   def _snake_case   (		self			):






     """simple docstring"""





     pass






   @unittest.skip(reason="Dinat does not use feedforward chunking"			)
   def _snake_case   (		self			):






     """simple docstring"""





     pass






   def _snake_case   (		self			):






     """simple docstring"""





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

     for model_class in self.all_model_classes:
       lowerCamelCase			: Union[str, Any]   =    model_class(__A			)
       self.assertIsInstance(model.get_input_embeddings()						,    (nn.Module)			)
       lowerCamelCase			: Tuple   =    model.get_output_embeddings()
       self.assertTrue(x is None or isinstance(__A						,    nn.Linear			)			)






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase   ,     lowerCamelCase			: str   =    self.model_tester.prepare_config_and_inputs_for_common()

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

       lowerCamelCase			: Any   =    ["pixel_values"]
       self.assertListEqual(arg_names[:1]						,    __A			)






   def _snake_case   (		self			):






     """simple docstring"""





     self.skipTest("Dinat's attention operation is handled entirely by NATTEN."			)






   def _snake_case   (		self						,    __A						,    __A						,    __A						,    __A			):






     """simple docstring"""





     lowerCamelCase			: Tuple   =    model_class(__A			)
     model.to(__A			)
     model.eval()

     with torch.no_grad():
       lowerCamelCase			: List[Any]   =    model(**self._prepare_for_class(__A						,    __A			)			)

     lowerCamelCase			: Any   =    outputs.hidden_states

     lowerCamelCase			: Tuple   =    getattr(
         self.model_tester						,    "expected_num_hidden_layers"						,    len(self.model_tester.depths			) + 1			)
     self.assertEqual(len(__A			)						,    __A			)

     # Dinat has a different seq_length
     lowerCamelCase			: List[str]   =    (
         config.patch_size
         if isinstance(config.patch_size						,    collections.abc.Iterable			)
         else (config.patch_size, config.patch_size)
     )

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

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






     """simple docstring"""





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

     lowerCamelCase			: Dict   =    (
         self.model_tester.image_size
         if isinstance(self.model_tester.image_size						,    collections.abc.Iterable			)
         else (self.model_tester.image_size, self.model_tester.image_size)
     )

     for model_class in self.all_model_classes:
       lowerCamelCase			: Union[str, Any]   =    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"]
       lowerCamelCase			: Tuple   =    True

       self.check_hidden_states_output(__A						,    __A						,    __A						,    __A			)






   @slow
   def _snake_case   (		self			):






     """simple docstring"""





     for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       lowerCamelCase			: Any   =    DinatModel.from_pretrained(__A			)
       self.assertIsNotNone(__A			)






   def _snake_case   (		self			):






     """simple docstring"""





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

     lowerCamelCase			: List[str]   =    _config_zero_init(__A			)
     for model_class in self.all_model_classes:
       lowerCamelCase			: List[Any]   =    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      UpperCAmelCase_ ( unittest.TestCase   ):

   '''simple docstring'''




   @cached_property
   def _snake_case   (		self			):






     """simple docstring"""





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






   @slow
   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Optional[Any]   =    DinatForImageClassification.from_pretrained("shi-labs/dinat-mini-in1k-224"			).to(__A			)
     lowerCamelCase			: int   =    self.default_image_processor

     lowerCamelCase			: Tuple   =    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"			)
     lowerCamelCase			: Union[str, Any]   =    image_processor(images=__A						,    return_tensors="pt"			).to(__A			)

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

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



@require_torch
@require_natten
class      UpperCAmelCase_ ( unittest.TestCase		,			UpperCamelCase   ):

   '''simple docstring'''




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






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Any   =    DinatModelTester(self			)







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


_snake_case					=					{
    '''configuration_rwkv''': ['''RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''RwkvConfig''', '''RwkvOnnxConfig'''],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _snake_case					=					[
       '''RWKV_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''RwkvForCausalLM''',
       '''RwkvModel''',
       '''RwkvPreTrainedModel''',
   ]


if TYPE_CHECKING:
   from .configuration_rwkv import RWKV_PRETRAINED_CONFIG_ARCHIVE_MAP, RwkvConfig, RwkvOnnxConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_rwkv import (
          RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
          RwkvForCausalLM,
          RwkvModel,
          RwkvPreTrainedModel,
      )
else:
   import sys

   _snake_case					=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



