code: 
  """simple docstring"""
  
  
  
  
  
  
  
  import math
  from dataclasses import dataclass
  from typing import List, Optional, Tuple, Union
  
  import numpy as np
  import torch
  
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..utils import BaseOutput, logging, randn_tensor
  from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
  
  
  SCREAMING_SNAKE_CASE_     				=   logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  
  
  
  
  @dataclass
  # Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EulerAncestralDiscrete
  class     __a			(					lowerCAmelCase__			):
         _lowerCamelCase							:							Union[str, Any]       =			42
         _lowerCamelCase							:							Tuple       =			None
  
  
  
  def 			__UpperCAmelCase		(		a_			:		Tuple    ,  a_			:		Union[str, Any]=0.999    ,  a_			:		Tuple="cosine"    ,  )     ->       Union[str, Any]:
     if alpha_transform_type == "cosine":
  
        def alpha_bar_fn(a_			:		Optional[Any]				):
           return math.cos((t + 0.008) / 1.008 * math.pi / 2				) ** 2
  
     elif alpha_transform_type == "exp":
  
        def alpha_bar_fn(a_			:		Tuple				):
           return math.exp(t * -12.0				)
  
     else:
        raise ValueError(F'''Unsupported alpha_tranform_type: {alpha_transform_type}'''				)
  
     _UpperCAmelCase   	=					[]
     for i in range(a_				):
        _UpperCAmelCase   	=					i / num_diffusion_timesteps
        _UpperCAmelCase   	=					(i + 1) / num_diffusion_timesteps
        betas.append(min(1 - alpha_bar_fn(a_				) / alpha_bar_fn(a_				)    ,  a_				)				)
     return torch.tensor(a_    ,  dtype=torch.floataa				)
  
  
  
  
  
  
  class     __a			(					lowerCAmelCase__				,	lowerCAmelCase__			):
         _lowerCamelCase							:							Any       =			[e.name for e in KarrasDiffusionSchedulers]
         _lowerCamelCase							:							Dict       =			1
  
  
  
  
  
  
  
         @register_to_config
         def __init__(					self							,lowerCamelCase_ = 1000							,lowerCamelCase_ = 0.0_0_0_1							,lowerCamelCase_ = 0.0_2							,lowerCamelCase_ = "linear"							,lowerCamelCase_ = None							,lowerCamelCase_ = "epsilon"							,lowerCamelCase_ = "linspace"							,lowerCamelCase_ = 0							,):
  
  
  
            '''simple docstring'''
  
            if trained_betas is not None:
               _UpperCAmelCase   	=					torch.tensor(_lowerCamelCase							,dtype=torch.floataa		)
            elif beta_schedule == "linear":
               _UpperCAmelCase   	=					torch.linspace(_lowerCamelCase							,_lowerCamelCase							,_lowerCamelCase							,dtype=torch.floataa		)
            elif beta_schedule == "scaled_linear":
               # this schedule is very specific to the latent diffusion model.
               _UpperCAmelCase   	=					(
                   torch.linspace(beta_start**0.5							,beta_end**0.5							,_lowerCamelCase							,dtype=torch.floataa		) ** 2
               )
            elif beta_schedule == "squaredcos_cap_v2":
               # Glide cosine schedule
               _UpperCAmelCase   	=					betas_for_alpha_bar(_lowerCamelCase		)
            else:
               raise NotImplementedError(f'''{beta_schedule} does is not implemented for {self.__class__}'''		)
  
            _UpperCAmelCase   	=					1.0 - self.betas
            _UpperCAmelCase   	=					torch.cumprod(self.alphas							,dim=0		)
  
            _UpperCAmelCase   	=					np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5		)
            _UpperCAmelCase   	=					np.concatenate([sigmas[::-1], [0.0]]		).astype(np.floataa		)
            _UpperCAmelCase   	=					torch.from_numpy(_lowerCamelCase		)
  
            # setable values
            _UpperCAmelCase   	=					None
            _UpperCAmelCase   	=					np.linspace(0							,num_train_timesteps - 1							,_lowerCamelCase							,dtype=_lowerCamelCase		)[::-1].copy()
            _UpperCAmelCase   	=					torch.from_numpy(_lowerCamelCase		)
            _UpperCAmelCase   	=					False
  
  
  
  
  
  
  
         @property
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            if self.config.timestep_spacing in ["linspace", "trailing"]:
               return self.sigmas.max()
  
            return (self.sigmas.max() ** 2 + 1) ** 0.5
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            if isinstance(_lowerCamelCase							,torch.Tensor		):
               _UpperCAmelCase   	=					timestep.to(self.timesteps.device		)
            _UpperCAmelCase   	=					(self.timesteps == timestep).nonzero().item()
            _UpperCAmelCase   	=					self.sigmas[step_index]
            _UpperCAmelCase   	=					sample / ((sigma**2 + 1) ** 0.5)
            _UpperCAmelCase   	=					True
            return sample
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_ = None		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					num_inference_steps
  
            # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
            if self.config.timestep_spacing == "linspace":
               _UpperCAmelCase   	=					np.linspace(0							,self.config.num_train_timesteps - 1							,_lowerCamelCase							,dtype=_lowerCamelCase		)[
                   ::-1
               ].copy()
            elif self.config.timestep_spacing == "leading":
               _UpperCAmelCase   	=					self.config.num_train_timesteps // self.num_inference_steps
               # creates integer timesteps by multiplying by ratio
               # casting to int to avoid issues when num_inference_step is power of 3
               _UpperCAmelCase   	=					(np.arange(0							,_lowerCamelCase		) * step_ratio).round()[::-1].copy().astype(_lowerCamelCase		)
               timesteps += self.config.steps_offset
            elif self.config.timestep_spacing == "trailing":
               _UpperCAmelCase   	=					self.config.num_train_timesteps / self.num_inference_steps
               # creates integer timesteps by multiplying by ratio
               # casting to int to avoid issues when num_inference_step is power of 3
               _UpperCAmelCase   	=					(np.arange(self.config.num_train_timesteps							,0							,-step_ratio		)).round().copy().astype(_lowerCamelCase		)
               timesteps -= 1
            else:
               raise ValueError(
                   f'''{self.config.timestep_spacing} is not supported. Please make sure to choose one of \'linspace\', \'leading\' or \'trailing\'.'''		)
  
            _UpperCAmelCase   	=					np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5		)
            _UpperCAmelCase   	=					np.interp(_lowerCamelCase							,np.arange(0							,len(_lowerCamelCase		)		)							,_lowerCamelCase		)
            _UpperCAmelCase   	=					np.concatenate([sigmas, [0.0]]		).astype(np.floataa		)
            _UpperCAmelCase   	=					torch.from_numpy(_lowerCamelCase		).to(device=_lowerCamelCase		)
            if str(_lowerCamelCase		).startswith("mps"		):
               # mps does not support float64
               _UpperCAmelCase   	=					torch.from_numpy(_lowerCamelCase		).to(_lowerCamelCase							,dtype=torch.floataa		)
            else:
               _UpperCAmelCase   	=					torch.from_numpy(_lowerCamelCase		).to(device=_lowerCamelCase		)
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_ = None							,lowerCamelCase_ = True							,):
  
  
  
            '''simple docstring'''
  
            if (
                isinstance(_lowerCamelCase							,_lowerCamelCase		)
                or isinstance(_lowerCamelCase							,torch.IntTensor		)
                or isinstance(_lowerCamelCase							,torch.LongTensor		)
            ):
               raise ValueError(
                   (
                       "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
                       " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
                       " one of the `scheduler.timesteps` as a timestep."
                   )							,)
  
            if not self.is_scale_input_called:
               logger.warning(
                   "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
                   "See `StableDiffusionPipeline` for a usage example."		)
  
            if isinstance(_lowerCamelCase							,torch.Tensor		):
               _UpperCAmelCase   	=					timestep.to(self.timesteps.device		)
  
            _UpperCAmelCase   	=					(self.timesteps == timestep).nonzero().item()
            _UpperCAmelCase   	=					self.sigmas[step_index]
  
            # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
            if self.config.prediction_type == "epsilon":
               _UpperCAmelCase   	=					sample - sigma * model_output
            elif self.config.prediction_type == "v_prediction":
               # * c_out + input * c_skip
               _UpperCAmelCase   	=					model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
            elif self.config.prediction_type == "sample":
               raise NotImplementedError("prediction_type not implemented yet: sample"		)
            else:
               raise ValueError(
                   f'''prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`'''		)
  
            _UpperCAmelCase   	=					self.sigmas[step_index]
            _UpperCAmelCase   	=					self.sigmas[step_index + 1]
            _UpperCAmelCase   	=					(sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
            _UpperCAmelCase   	=					(sigma_to**2 - sigma_up**2) ** 0.5
  
            # 2. Convert to an ODE derivative
            _UpperCAmelCase   	=					(sample - pred_original_sample) / sigma
  
            _UpperCAmelCase   	=					sigma_down - sigma
  
            _UpperCAmelCase   	=					sample + derivative * dt
  
            _UpperCAmelCase   	=					model_output.device
            _UpperCAmelCase   	=					randn_tensor(model_output.shape							,dtype=model_output.dtype							,device=_lowerCamelCase							,generator=_lowerCamelCase		)
  
            _UpperCAmelCase   	=					prev_sample + noise * sigma_up
  
            if not return_dict:
               return (prev_sample,)
  
            return EulerAncestralDiscreteSchedulerOutput(
                prev_sample=_lowerCamelCase							,pred_original_sample=_lowerCamelCase		)
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.sigmas.to(device=original_samples.device							,dtype=original_samples.dtype		)
            if original_samples.device.type == "mps" and torch.is_floating_point(_lowerCamelCase		):
               # mps does not support float64
               _UpperCAmelCase   	=					self.timesteps.to(original_samples.device							,dtype=torch.floataa		)
               _UpperCAmelCase   	=					timesteps.to(original_samples.device							,dtype=torch.floataa		)
            else:
               _UpperCAmelCase   	=					self.timesteps.to(original_samples.device		)
               _UpperCAmelCase   	=					timesteps.to(original_samples.device		)
  
            _UpperCAmelCase   	=					[(schedule_timesteps == t).nonzero().item() for t in timesteps]
  
            _UpperCAmelCase   	=					sigmas[step_indices].flatten()
            while len(sigma.shape		) < len(original_samples.shape		):
               _UpperCAmelCase   	=					sigma.unsqueeze(-1		)
  
            _UpperCAmelCase   	=					original_samples + noise * sigma
            return noisy_samples
  
  
  
  
  
  
  
         def __len__(					self		):
  
  
  
            '''simple docstring'''
  
            return self.config.num_train_timesteps
  
  
style_context: 
  """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     __a			:
  
  
  
  
  
  
  
         def __init__(					self							,lowerCamelCase_							,lowerCamelCase_=13							,lowerCamelCase_=64							,lowerCamelCase_=4							,lowerCamelCase_=3							,lowerCamelCase_=16							,lowerCamelCase_=[1, 2, 1]							,lowerCamelCase_=[2, 4, 8]							,lowerCamelCase_=3							,lowerCamelCase_=2.0							,lowerCamelCase_=True							,lowerCamelCase_=0.0							,lowerCamelCase_=0.0							,lowerCamelCase_=0.1							,lowerCamelCase_="gelu"							,lowerCamelCase_=True							,lowerCamelCase_=0.0_2							,lowerCamelCase_=1e-5							,lowerCamelCase_=True							,lowerCamelCase_=None							,lowerCamelCase_=True							,lowerCamelCase_=10							,lowerCamelCase_=["stage1", "stage2"]							,lowerCamelCase_=[1, 2]							,):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					parent
            _UpperCAmelCase   	=					batch_size
            _UpperCAmelCase   	=					image_size
            _UpperCAmelCase   	=					patch_size
            _UpperCAmelCase   	=					num_channels
            _UpperCAmelCase   	=					embed_dim
            _UpperCAmelCase   	=					depths
            _UpperCAmelCase   	=					num_heads
            _UpperCAmelCase   	=					kernel_size
            _UpperCAmelCase   	=					mlp_ratio
            _UpperCAmelCase   	=					qkv_bias
            _UpperCAmelCase   	=					hidden_dropout_prob
            _UpperCAmelCase   	=					attention_probs_dropout_prob
            _UpperCAmelCase   	=					drop_path_rate
            _UpperCAmelCase   	=					hidden_act
            _UpperCAmelCase   	=					patch_norm
            _UpperCAmelCase   	=					layer_norm_eps
            _UpperCAmelCase   	=					initializer_range
            _UpperCAmelCase   	=					is_training
            _UpperCAmelCase   	=					scope
            _UpperCAmelCase   	=					use_labels
            _UpperCAmelCase   	=					num_labels
            _UpperCAmelCase   	=					out_features
            _UpperCAmelCase   	=					out_indices
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)
  
            _UpperCAmelCase   	=					None
            if self.use_labels:
               _UpperCAmelCase   	=					ids_tensor([self.batch_size]							,self.num_labels		)
  
            _UpperCAmelCase   	=					self.get_config()
  
            return config, pixel_values, labels
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''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    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					NatModel(config=lowerCamelCase_		)
            model.to(lowerCamelCase_		)
            model.eval()
            _UpperCAmelCase   	=					model(lowerCamelCase_		)
  
            _UpperCAmelCase   	=					_UpperCAmelCase   	=					(config.image_size // config.patch_size) // (2 ** (len(config.depths		) - 1))
            _UpperCAmelCase   	=					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							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					NatForImageClassification(lowerCamelCase_		)
            model.to(lowerCamelCase_		)
            model.eval()
            _UpperCAmelCase   	=					model(lowerCamelCase_							,labels=lowerCamelCase_		)
            self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_labels)		)
  
            # test greyscale images
            _UpperCAmelCase   	=					1
            _UpperCAmelCase   	=					NatForImageClassification(lowerCamelCase_		)
            model.to(lowerCamelCase_		)
            model.eval()
  
            _UpperCAmelCase   	=					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
            _UpperCAmelCase   	=					model(lowerCamelCase_		)
            self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_labels)		)
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					NatBackbone(config=lowerCamelCase_		)
            model.to(lowerCamelCase_		)
            model.eval()
            _UpperCAmelCase   	=					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
            _UpperCAmelCase   	=					None
            _UpperCAmelCase   	=					NatBackbone(config=lowerCamelCase_		)
            model.to(lowerCamelCase_		)
            model.eval()
            _UpperCAmelCase   	=					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    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.prepare_config_and_inputs()
            _UpperCAmelCase	,							_UpperCAmelCase	,							_UpperCAmelCase   	=					config_and_inputs
            _UpperCAmelCase   	=					{"pixel_values": pixel_values}
            return config, inputs_dict
  
  
  
  
  
  
  @require_natten
  @require_torch
  class     __a			(					lowerCAmelCase_				,	lowerCAmelCase_				,	unittest.TestCase			):
         _lowerCamelCase							:							str       =			(
             (
                 NatModel,
                 NatForImageClassification,
                 NatBackbone,
             )
             if is_torch_available()
             else ()
         )
         _lowerCamelCase							:							int       =			(
             {'feature-extraction': NatModel, 'image-classification': NatForImageClassification}
             if is_torch_available()
             else {}
         )
         _lowerCamelCase							:							Optional[Any]       =			False
  
         _lowerCamelCase							:							Optional[Any]       =			False
         _lowerCamelCase							:							Optional[int]       =			False
         _lowerCamelCase							:							List[str]       =			False
         _lowerCamelCase							:							str       =			False
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					NatModelTester(self		)
            _UpperCAmelCase   	=					ConfigTester(self							,config_class=lowerCamelCase_							,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'''
  
            _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*lowerCamelCase_		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_image_classification(*lowerCamelCase_		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					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    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            pass
  
  
  
  
  
  
  
         @unittest.skip(reason="Nat does not use feedforward chunking"		)
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            pass
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase	,							_UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs_for_common()
  
            for model_class in self.all_model_classes:
               _UpperCAmelCase   	=					model_class(lowerCamelCase_		)
               self.assertIsInstance(model.get_input_embeddings()							,(nn.Module)		)
               _UpperCAmelCase   	=					model.get_output_embeddings()
               self.assertTrue(x is None or isinstance(lowerCamelCase_							,nn.Linear		)		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase	,							_UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs_for_common()
  
            for model_class in self.all_model_classes:
               _UpperCAmelCase   	=					model_class(lowerCamelCase_		)
               _UpperCAmelCase   	=					inspect.signature(model.forward		)
               # signature.parameters is an OrderedDict => so arg_names order is deterministic
               _UpperCAmelCase   	=					[*signature.parameters.keys()]
  
               _UpperCAmelCase   	=					["pixel_values"]
               self.assertListEqual(arg_names[:1]							,lowerCamelCase_		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            self.skipTest("Nat's attention operation is handled entirely by NATTEN."		)
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					model_class(lowerCamelCase_		)
            model.to(lowerCamelCase_		)
            model.eval()
  
            with torch.no_grad():
               _UpperCAmelCase   	=					model(**self._prepare_for_class(lowerCamelCase_							,lowerCamelCase_		)		)
  
            _UpperCAmelCase   	=					outputs.hidden_states
  
            _UpperCAmelCase   	=					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
            _UpperCAmelCase   	=					(
                config.patch_size
                if isinstance(config.patch_size							,collections.abc.Iterable		)
                else (config.patch_size, config.patch_size)
            )
  
            _UpperCAmelCase   	=					image_size[0] // patch_size[0]
            _UpperCAmelCase   	=					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":
               _UpperCAmelCase   	=					outputs.reshaped_hidden_states
               self.assertEqual(len(lowerCamelCase_		)							,lowerCamelCase_		)
  
               _UpperCAmelCase	,							_UpperCAmelCase	,							_UpperCAmelCase	,							_UpperCAmelCase   	=					reshaped_hidden_states[0].shape
               _UpperCAmelCase   	=					(
                   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    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase	,							_UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs_for_common()
  
            _UpperCAmelCase   	=					(
                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:
               _UpperCAmelCase   	=					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"]
               _UpperCAmelCase   	=					True
  
               self.check_hidden_states_output(lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		)
  
  
  
  
  
  
  
         @slow
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _UpperCAmelCase   	=					NatModel.from_pretrained(lowerCamelCase_		)
               self.assertIsNotNone(lowerCamelCase_		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase	,							_UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs_for_common()
  
            _UpperCAmelCase   	=					_config_zero_init(lowerCamelCase_		)
            for model_class in self.all_model_classes:
               _UpperCAmelCase   	=					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			):
  
  
  
  
  
  
  
         @cached_property
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            return AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224"		) if is_vision_available() else None
  
  
  
  
  
  
  
         @slow
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					NatForImageClassification.from_pretrained("shi-labs/nat-mini-in1k-224"		).to(lowerCamelCase_		)
            _UpperCAmelCase   	=					self.default_image_processor
  
            _UpperCAmelCase   	=					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"		)
            _UpperCAmelCase   	=					image_processor(images=lowerCamelCase_							,return_tensors="pt"		).to(lowerCamelCase_		)
  
            # forward pass
            with torch.no_grad():
               _UpperCAmelCase   	=					model(**lowerCamelCase_		)
  
            # verify the logits
            _UpperCAmelCase   	=					torch.Size((1, 1000)		)
            self.assertEqual(outputs.logits.shape							,lowerCamelCase_		)
            _UpperCAmelCase   	=					torch.tensor([0.3_8_0_5, -0.8_6_7_6, -0.3_9_1_2]		).to(lowerCamelCase_		)
            self.assertTrue(torch.allclose(outputs.logits[0, :3]							,lowerCamelCase_							,atol=1e-4		)		)
  
  
  
  
  
  
  @require_torch
  @require_natten
  class     __a			(					unittest.TestCase				,	lowerCAmelCase_			):
         _lowerCamelCase							:							Any       =			(NatBackbone,) if is_torch_available() else ()
         _lowerCamelCase							:							Dict       =			NatConfig
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					NatModelTester(self		)
  
  