
'''simple docstring'''







from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block


@dataclass
class 						__UpperCAmelCase      (					_lowerCamelCase						):
							__lowercase      		=       42


class 						__UpperCAmelCase      (					nn.Module						):


							def __init__(       self    ,   lowerCAmelCase_=3    ,   lowerCAmelCase_=3    ,   lowerCAmelCase_=("DownEncoderBlock2D",)    ,   lowerCAmelCase_=(64,)    ,   lowerCAmelCase_=2    ,   lowerCAmelCase_=32    ,   lowerCAmelCase_="silu"    ,   lowerCAmelCase_=True    ,   ):





								"""simple docstring"""






								super().__init__()
								_snake_case   						=		layers_per_block

								_snake_case   						=		torch.nn.Convad(
								    lowerCAmelCase_    ,   block_out_channels[0]    ,   kernel_size=3    ,   stride=1    ,   padding=1    ,   )

								_snake_case   						=		None
								_snake_case   						=		nn.ModuleList([]  )

								# down
								_snake_case   						=		block_out_channels[0]
								for i, down_block_type in enumerate(lowerCAmelCase_  ):
									_snake_case   						=		output_channel
									_snake_case   						=		block_out_channels[i]
									_snake_case   						=		i == len(lowerCAmelCase_  ) - 1

									_snake_case   						=		get_down_block(
									    lowerCAmelCase_    ,   num_layers=self.layers_per_block    ,   in_channels=lowerCAmelCase_    ,   out_channels=lowerCAmelCase_    ,   add_downsample=not is_final_block    ,   resnet_eps=1E-6    ,   downsample_padding=0    ,   resnet_act_fn=lowerCAmelCase_    ,   resnet_groups=lowerCAmelCase_    ,   attention_head_dim=lowerCAmelCase_    ,   temb_channels=lowerCAmelCase_    ,   )
									self.down_blocks.append(lowerCAmelCase_  )

								# mid
								_snake_case   						=		UNetMidBlockaD(
								    in_channels=block_out_channels[-1]    ,   resnet_eps=1E-6    ,   resnet_act_fn=lowerCAmelCase_    ,   output_scale_factor=1    ,   resnet_time_scale_shift='default'    ,   attention_head_dim=block_out_channels[-1]    ,   resnet_groups=lowerCAmelCase_    ,   temb_channels=lowerCAmelCase_    ,   )

								# out
								_snake_case   						=		nn.GroupNorm(num_channels=block_out_channels[-1]    ,   num_groups=lowerCAmelCase_    ,   eps=1E-6  )
								_snake_case   						=		nn.SiLU()

								_snake_case   						=		2 * out_channels if double_z else out_channels
								_snake_case   						=		nn.Convad(block_out_channels[-1]    ,   lowerCAmelCase_    ,   3    ,   padding=1  )

								_snake_case   						=		False




							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		x
								_snake_case   						=		self.conv_in(lowerCAmelCase_  )

								if self.training and self.gradient_checkpointing:

									def create_custom_forward(lowerCAmelCase_  ):
										def custom_forward(*lowerCAmelCase_  ):
											return module(*lowerCAmelCase_  )

										return custom_forward

									# down
									if is_torch_version('>='    ,   '1.11.0'  ):
										for down_block in self.down_blocks:
											_snake_case   						=		torch.utils.checkpoint.checkpoint(
											    create_custom_forward(lowerCAmelCase_  )    ,   lowerCAmelCase_    ,   use_reentrant=lowerCAmelCase_  )
										# middle
										_snake_case   						=		torch.utils.checkpoint.checkpoint(
										    create_custom_forward(self.mid_block  )    ,   lowerCAmelCase_    ,   use_reentrant=lowerCAmelCase_  )
									else:
										for down_block in self.down_blocks:
											_snake_case   						=		torch.utils.checkpoint.checkpoint(create_custom_forward(lowerCAmelCase_  )    ,   lowerCAmelCase_  )
										# middle
										_snake_case   						=		torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block  )    ,   lowerCAmelCase_  )

								else:
									# down
									for down_block in self.down_blocks:
										_snake_case   						=		down_block(lowerCAmelCase_  )

									# middle
									_snake_case   						=		self.mid_block(lowerCAmelCase_  )

								# post-process
								_snake_case   						=		self.conv_norm_out(lowerCAmelCase_  )
								_snake_case   						=		self.conv_act(lowerCAmelCase_  )
								_snake_case   						=		self.conv_out(lowerCAmelCase_  )

								return sample


class 						__UpperCAmelCase      (					nn.Module						):


							def __init__(       self    ,   lowerCAmelCase_=3    ,   lowerCAmelCase_=3    ,   lowerCAmelCase_=("UpDecoderBlock2D",)    ,   lowerCAmelCase_=(64,)    ,   lowerCAmelCase_=2    ,   lowerCAmelCase_=32    ,   lowerCAmelCase_="silu"    ,   lowerCAmelCase_="group"    ,   ):





								"""simple docstring"""






								super().__init__()
								_snake_case   						=		layers_per_block

								_snake_case   						=		nn.Convad(
								    lowerCAmelCase_    ,   block_out_channels[-1]    ,   kernel_size=3    ,   stride=1    ,   padding=1    ,   )

								_snake_case   						=		None
								_snake_case   						=		nn.ModuleList([]  )

								_snake_case   						=		in_channels if norm_type == 'spatial' else None

								# mid
								_snake_case   						=		UNetMidBlockaD(
								    in_channels=block_out_channels[-1]    ,   resnet_eps=1E-6    ,   resnet_act_fn=lowerCAmelCase_    ,   output_scale_factor=1    ,   resnet_time_scale_shift='default' if norm_type == 'group' else norm_type    ,   attention_head_dim=block_out_channels[-1]    ,   resnet_groups=lowerCAmelCase_    ,   temb_channels=lowerCAmelCase_    ,   )

								# up
								_snake_case   						=		list(reversed(lowerCAmelCase_  )  )
								_snake_case   						=		reversed_block_out_channels[0]
								for i, up_block_type in enumerate(lowerCAmelCase_  ):
									_snake_case   						=		output_channel
									_snake_case   						=		reversed_block_out_channels[i]

									_snake_case   						=		i == len(lowerCAmelCase_  ) - 1

									_snake_case   						=		get_up_block(
									    lowerCAmelCase_    ,   num_layers=self.layers_per_block + 1    ,   in_channels=lowerCAmelCase_    ,   out_channels=lowerCAmelCase_    ,   prev_output_channel=lowerCAmelCase_    ,   add_upsample=not is_final_block    ,   resnet_eps=1E-6    ,   resnet_act_fn=lowerCAmelCase_    ,   resnet_groups=lowerCAmelCase_    ,   attention_head_dim=lowerCAmelCase_    ,   temb_channels=lowerCAmelCase_    ,   resnet_time_scale_shift=lowerCAmelCase_    ,   )
									self.up_blocks.append(lowerCAmelCase_  )
									_snake_case   						=		output_channel

								# out
								if norm_type == "spatial":
									_snake_case   						=		SpatialNorm(block_out_channels[0]    ,   lowerCAmelCase_  )
								else:
									_snake_case   						=		nn.GroupNorm(num_channels=block_out_channels[0]    ,   num_groups=lowerCAmelCase_    ,   eps=1E-6  )
								_snake_case   						=		nn.SiLU()
								_snake_case   						=		nn.Convad(block_out_channels[0]    ,   lowerCAmelCase_    ,   3    ,   padding=1  )

								_snake_case   						=		False




							def  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_=None  ):





								"""simple docstring"""






								_snake_case   						=		z
								_snake_case   						=		self.conv_in(lowerCAmelCase_  )

								_snake_case   						=		next(iter(self.up_blocks.parameters()  )  ).dtype
								if self.training and self.gradient_checkpointing:

									def create_custom_forward(lowerCAmelCase_  ):
										def custom_forward(*lowerCAmelCase_  ):
											return module(*lowerCAmelCase_  )

										return custom_forward

									if is_torch_version('>='    ,   '1.11.0'  ):
										# middle
										_snake_case   						=		torch.utils.checkpoint.checkpoint(
										    create_custom_forward(self.mid_block  )    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   use_reentrant=lowerCAmelCase_  )
										_snake_case   						=		sample.to(lowerCAmelCase_  )

										# up
										for up_block in self.up_blocks:
											_snake_case   						=		torch.utils.checkpoint.checkpoint(
											    create_custom_forward(lowerCAmelCase_  )    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   use_reentrant=lowerCAmelCase_  )
									else:
										# middle
										_snake_case   						=		torch.utils.checkpoint.checkpoint(
										    create_custom_forward(self.mid_block  )    ,   lowerCAmelCase_    ,   lowerCAmelCase_  )
										_snake_case   						=		sample.to(lowerCAmelCase_  )

										# up
										for up_block in self.up_blocks:
											_snake_case   						=		torch.utils.checkpoint.checkpoint(create_custom_forward(lowerCAmelCase_  )    ,   lowerCAmelCase_    ,   lowerCAmelCase_  )
								else:
									# middle
									_snake_case   						=		self.mid_block(lowerCAmelCase_    ,   lowerCAmelCase_  )
									_snake_case   						=		sample.to(lowerCAmelCase_  )

									# up
									for up_block in self.up_blocks:
										_snake_case   						=		up_block(lowerCAmelCase_    ,   lowerCAmelCase_  )

        # post-process
								if latent_embeds is None:
									_snake_case   						=		self.conv_norm_out(lowerCAmelCase_  )
								else:
									_snake_case   						=		self.conv_norm_out(lowerCAmelCase_    ,   lowerCAmelCase_  )
								_snake_case   						=		self.conv_act(lowerCAmelCase_  )
								_snake_case   						=		self.conv_out(lowerCAmelCase_  )

								return sample


class 						__UpperCAmelCase      (					nn.Module						):


							def __init__(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_=None    ,   lowerCAmelCase_="random"    ,   lowerCAmelCase_=False    ,   lowerCAmelCase_=True  ):





								"""simple docstring"""






								super().__init__()
								_snake_case   						=		n_e
								_snake_case   						=		vq_embed_dim
								_snake_case   						=		beta
								_snake_case   						=		legacy

								_snake_case   						=		nn.Embedding(self.n_e    ,   self.vq_embed_dim  )
								self.embedding.weight.data.uniform_(-1.0 / self.n_e    ,   1.0 / self.n_e  )

								_snake_case   						=		remap
								if self.remap is not None:
									self.register_buffer('used'    ,   torch.tensor(np.load(self.remap  )  )  )
									_snake_case   						=		self.used.shape[0]
									_snake_case   						=		unknown_index  # "random" or "extra" or integer
									if self.unknown_index == "extra":
										_snake_case   						=		self.re_embed
										_snake_case   						=		self.re_embed + 1
									print(
									    F'Remapping {self.n_e} indices to {self.re_embed} indices. '
									    F'Using {self.unknown_index} for unknown indices.'  )
								else:
									_snake_case   						=		n_e

								_snake_case   						=		sane_index_shape


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		inds.shape
								assert len(lowerCAmelCase_  ) > 1
								_snake_case   						=		inds.reshape(ishape[0]    ,   -1  )
								_snake_case   						=		self.used.to(lowerCAmelCase_  )
								_snake_case   						=		(inds[:, :, None] == used[None, None, ...]).long()
								_snake_case   						=		match.argmax(-1  )
								_snake_case   						=		match.sum(2  ) < 1
								if self.unknown_index == "random":
									_snake_case   						=		torch.randint(0    ,   self.re_embed    ,   size=new[unknown].shape  ).to(device=new.device  )
								else:
									_snake_case   						=		self.unknown_index
								return new.reshape(lowerCAmelCase_  )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		inds.shape
								assert len(lowerCAmelCase_  ) > 1
								_snake_case   						=		inds.reshape(ishape[0]    ,   -1  )
								_snake_case   						=		self.used.to(lowerCAmelCase_  )
								if self.re_embed > self.used.shape[0]:  # extra token
									_snake_case   						=		0  # simply set to zero
								_snake_case   						=		torch.gather(used[None, :][inds.shape[0] * [0], :]    ,   1    ,   lowerCAmelCase_  )
								return back.reshape(lowerCAmelCase_  )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		z.permute(0    ,   2    ,   3    ,   1  ).contiguous()
								_snake_case   						=		z.view(-1    ,   self.vq_embed_dim  )

								# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
								_snake_case   						=		torch.argmin(torch.cdist(lowerCAmelCase_    ,   self.embedding.weight  )    ,   dim=1  )

								_snake_case   						=		self.embedding(lowerCAmelCase_  ).view(z.shape  )
								_snake_case   						=		None
								_snake_case   						=		None

								# compute loss for embedding
								if not self.legacy:
									_snake_case   						=		self.beta * torch.mean((z_q.detach() - z) ** 2  ) + torch.mean((z_q - z.detach()) ** 2  )
								else:
									_snake_case   						=		torch.mean((z_q.detach() - z) ** 2  ) + self.beta * torch.mean((z_q - z.detach()) ** 2  )

								# preserve gradients
								_snake_case   						=		z + (z_q - z).detach()

								# reshape back to match original input shape
								_snake_case   						=		z_q.permute(0    ,   3    ,   1    ,   2  ).contiguous()

								if self.remap is not None:
									_snake_case   						=		min_encoding_indices.reshape(z.shape[0]    ,   -1  )  # add batch axis
									_snake_case   						=		self.remap_to_used(lowerCAmelCase_  )
									_snake_case   						=		min_encoding_indices.reshape(-1    ,   1  )  # flatten

								if self.sane_index_shape:
									_snake_case   						=		min_encoding_indices.reshape(z_q.shape[0]    ,   z_q.shape[2]    ,   z_q.shape[3]  )

								return z_q, loss, (perplexity, min_encodings, min_encoding_indices)




							def  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								if self.remap is not None:
									_snake_case   						=		indices.reshape(shape[0]    ,   -1  )  # add batch axis
									_snake_case   						=		self.unmap_to_all(lowerCAmelCase_  )
									_snake_case   						=		indices.reshape(-1  )  # flatten again

								# get quantized latent vectors
								_snake_case   						=		self.embedding(lowerCAmelCase_  )

								if shape is not None:
									_snake_case   						=		z_q.view(lowerCAmelCase_  )
									# reshape back to match original input shape
									_snake_case   						=		z_q.permute(0    ,   3    ,   1    ,   2  ).contiguous()

								return z_q


class 						__UpperCAmelCase      (					_lowerCamelCase						):


							def __init__(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_=False  ):





								"""simple docstring"""






								_snake_case   						=		parameters
								_snake_case						, _snake_case   						=		torch.chunk(lowerCAmelCase_    ,   2    ,   dim=1  )
								_snake_case   						=		torch.clamp(self.logvar    ,   -30.0    ,   20.0  )
								_snake_case   						=		deterministic
								_snake_case   						=		torch.exp(0.5 * self.logvar  )
								_snake_case   						=		torch.exp(self.logvar  )
								if self.deterministic:
									_snake_case   						=		_snake_case   						=		torch.zeros_like(
									    self.mean    ,   device=self.parameters.device    ,   dtype=self.parameters.dtype  )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_ = None  ):





								"""simple docstring"""






								_snake_case   						=		randn_tensor(
								    self.mean.shape    ,   generator=lowerCAmelCase_    ,   device=self.parameters.device    ,   dtype=self.parameters.dtype  )
								_snake_case   						=		self.mean + self.std * sample
								return x


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_=None  ):





								"""simple docstring"""






								if self.deterministic:
									return torch.Tensor([0.0]  )
								else:
									if other is None:
										return 0.5 * torch.sum(torch.pow(self.mean    ,   2  ) + self.var - 1.0 - self.logvar    ,   dim=[1, 2, 3]  )
									else:
										return 0.5 * torch.sum(
										    torch.pow(self.mean - other.mean    ,   2  ) / other.var
										    + self.var / other.var
										    - 1.0
										    - self.logvar
										    + other.logvar    ,   dim=[1, 2, 3]    ,   )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_=[1, 2, 3]  ):





								"""simple docstring"""






								if self.deterministic:
									return torch.Tensor([0.0]  )
								_snake_case   						=		np.log(2.0 * np.pi  )
								return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean    ,   2  ) / self.var    ,   dim=lowerCAmelCase_  )




							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								return self.mean


"""simple docstring"""







import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
           MODEL_MAPPING,
           LevitForImageClassification,
           LevitForImageClassificationWithTeacher,
           LevitModel,
       )
       from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       from PIL import Image

       from transformers import LevitImageProcessor







class        UpperCamelCase		(			lowercase_  ):
   def 			_UpperCAmelCase   (			self			)	->      Any:


        '''simple docstring'''


        lowercase_	:				Any													=	self.config_class(**self.inputs_dict			)
        self.parent.assertTrue(hasattr(__UpperCamelCase     ,'hidden_sizes'			)			)
        self.parent.assertTrue(hasattr(__UpperCamelCase     ,'num_attention_heads'			)			)







class        UpperCamelCase		:
   def __init__(			self     ,__UpperCamelCase     ,__UpperCamelCase=13     ,__UpperCamelCase=64     ,__UpperCamelCase=3     ,__UpperCamelCase=3     ,__UpperCamelCase=2     ,__UpperCamelCase=1     ,__UpperCamelCase=16     ,__UpperCamelCase=[128, 256, 384]     ,__UpperCamelCase=[4, 6, 8]     ,__UpperCamelCase=[2, 3, 4]     ,__UpperCamelCase=[16, 16, 16]     ,__UpperCamelCase=0     ,__UpperCamelCase=[2, 2, 2]     ,__UpperCamelCase=[2, 2, 2]     ,__UpperCamelCase=0.02     ,__UpperCamelCase=True     ,__UpperCamelCase=True     ,__UpperCamelCase=2     ,)	->      str:


        '''simple docstring'''


        lowercase_	:				Optional[Any]													=	parent
        lowercase_	:				Any													=	batch_size
        lowercase_	:				str													=	image_size
        lowercase_	:				List[Any]													=	num_channels
        lowercase_	:				Any													=	kernel_size
        lowercase_	:				Any													=	stride
        lowercase_	:				List[str]													=	padding
        lowercase_	:				Tuple													=	hidden_sizes
        lowercase_	:				List[Any]													=	num_attention_heads
        lowercase_	:				Optional[int]													=	depths
        lowercase_	:				Union[str, Any]													=	key_dim
        lowercase_	:				List[str]													=	drop_path_rate
        lowercase_	:				Dict													=	patch_size
        lowercase_	:				List[Any]													=	attention_ratio
        lowercase_	:				Optional[int]													=	mlp_ratio
        lowercase_	:				Any													=	initializer_range
        lowercase_	:				Optional[Any]													=	[
            ['Subsample', key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
            ['Subsample', key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
        ]
        lowercase_	:				Dict													=	is_training
        lowercase_	:				Tuple													=	use_labels
        lowercase_	:				Dict													=	num_labels
        lowercase_	:				Union[str, Any]													=	initializer_range
   def 			_UpperCAmelCase   (			self			)	->      List[str]:


        '''simple docstring'''


        lowercase_	:				Union[str, Any]													=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

        lowercase_	:				List[str]													=	None
        if self.use_labels:
             lowercase_	:				List[Any]													=	ids_tensor([self.batch_size]     ,self.num_labels			)

        lowercase_	:				List[str]													=	self.get_config()
        return config, pixel_values, labels
   def 			_UpperCAmelCase   (			self			)	->      List[str]:


        '''simple docstring'''


        return LevitConfig(
            image_size=self.image_size     ,num_channels=self.num_channels     ,kernel_size=self.kernel_size     ,stride=self.stride     ,padding=self.padding     ,patch_size=self.patch_size     ,hidden_sizes=self.hidden_sizes     ,num_attention_heads=self.num_attention_heads     ,depths=self.depths     ,key_dim=self.key_dim     ,drop_path_rate=self.drop_path_rate     ,mlp_ratio=self.mlp_ratio     ,attention_ratio=self.attention_ratio     ,initializer_range=self.initializer_range     ,down_ops=self.down_ops     ,)
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)	->      Any:


        '''simple docstring'''


        lowercase_	:				Union[str, Any]													=	LevitModel(config=__UpperCamelCase			)
        model.to(__UpperCamelCase			)
        model.eval()
        lowercase_	:				Optional[Any]													=	model(__UpperCamelCase			)
        lowercase_	:				Union[str, Any]													=	(self.image_size, self.image_size)
        lowercase_							,  lowercase_	:				Tuple													=	image_size[0], image_size[1]
        for _ in range(4			):
             lowercase_	:				Optional[int]													=	floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1			)
             lowercase_	:				Dict													=	floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1			)
        self.parent.assertEqual(
            result.last_hidden_state.shape     ,(self.batch_size, ceil(height / 4			) * ceil(width / 4			), self.hidden_sizes[-1])     ,)
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)	->      Union[str, Any]:


        '''simple docstring'''


        lowercase_	:				str													=	self.num_labels
        lowercase_	:				Optional[Any]													=	LevitForImageClassification(__UpperCamelCase			)
        model.to(__UpperCamelCase			)
        model.eval()
        lowercase_	:				Union[str, Any]													=	model(__UpperCamelCase     ,labels=__UpperCamelCase			)
        self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.num_labels)			)







   def 			_UpperCAmelCase   (			self			)	->      List[Any]:


        '''simple docstring'''


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







@require_torch
class        UpperCamelCase		(			lowercase_						,      lowercase_						,      unittest.TestCase  ):
   lowercase							=  (
       (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
       if is_torch_available()
       else ()
   )
   lowercase							=  (
       {
           'feature-extraction': LevitModel,
           'image-classification': (LevitForImageClassification, LevitForImageClassificationWithTeacher),
       }
       if is_torch_available()
       else {}
   )

   lowercase							=  False
   lowercase							=  False
   lowercase							=  False
   lowercase							=  False
   lowercase							=  False
   def 			_UpperCAmelCase   (			self			)	->      List[str]:


        '''simple docstring'''


        lowercase_	:				Tuple													=	LevitModelTester(self			)
        lowercase_	:				Dict													=	ConfigTester(self     ,config_class=__UpperCamelCase     ,has_text_modality=__UpperCamelCase     ,hidden_size=37			)
   def 			_UpperCAmelCase   (			self			)	->      List[str]:


        '''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			)	->      List[Any]:


        '''simple docstring'''


        return
   @unittest.skip(reason='Levit does not use inputs_embeds'			)
   def 			_UpperCAmelCase   (			self			)	->      Tuple:


        '''simple docstring'''


        pass
   @unittest.skip(reason='Levit does not support input and output embeddings'			)
   def 			_UpperCAmelCase   (			self			)	->      Tuple:


        '''simple docstring'''


        pass
   @unittest.skip(reason='Levit does not output attentions'			)
   def 			_UpperCAmelCase   (			self			)	->      Optional[Any]:


        '''simple docstring'''


        pass
   def 			_UpperCAmelCase   (			self			)	->      List[str]:


        '''simple docstring'''


        lowercase_							,  lowercase_	:				Optional[int]													=	self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
             lowercase_	:				Optional[int]													=	model_class(__UpperCamelCase			)
             lowercase_	:				Dict													=	inspect.signature(model.forward			)
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             lowercase_	:				Union[str, Any]													=	[*signature.parameters.keys()]

             lowercase_	:				Optional[int]													=	['pixel_values']
             self.assertListEqual(arg_names[:1]     ,__UpperCamelCase			)
   def 			_UpperCAmelCase   (			self			)	->      Dict:


        '''simple docstring'''


        def check_hidden_states_output(__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			):
             lowercase_	:				List[Any]													=	model_class(__UpperCamelCase			)
             model.to(__UpperCamelCase			)
             model.eval()

             with torch.no_grad():
                  lowercase_	:				Optional[int]													=	model(**self._prepare_for_class(__UpperCamelCase     ,__UpperCamelCase			)			)

             lowercase_	:				Dict													=	outputs.hidden_states

             lowercase_	:				List[str]													=	len(self.model_tester.depths			) + 1
             self.assertEqual(len(__UpperCamelCase			)     ,__UpperCamelCase			)

             lowercase_	:				Optional[int]													=	(self.model_tester.image_size, self.model_tester.image_size)
             lowercase_							,  lowercase_	:				Union[str, Any]													=	image_size[0], image_size[1]
             for _ in range(4			):
                  lowercase_	:				Optional[Any]													=	floor(
                      (
                          (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                          / self.model_tester.stride
                      )
                      + 1			)
                  lowercase_	:				Optional[int]													=	floor(
                      (
                          (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                          / self.model_tester.stride
                      )
                      + 1			)
             # verify the first hidden states (first block)
             self.assertListEqual(
                 list(hidden_states[0].shape[-2:]			)     ,[
                     height * width,
                     self.model_tester.hidden_sizes[0],
                 ]     ,)

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

        for model_class in self.all_model_classes:
             lowercase_	:				Any													=	True
             check_hidden_states_output(__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             lowercase_	:				Optional[int]													=	True

             check_hidden_states_output(__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)
   @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'			)
   def 			_UpperCAmelCase   (			self			)	->      Tuple:


        '''simple docstring'''


        pass
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase=False			)	->      List[Any]:


        '''simple docstring'''


        lowercase_	:				Optional[Any]													=	super()._prepare_for_class(__UpperCamelCase     ,__UpperCamelCase     ,return_labels=__UpperCamelCase			)

        if return_labels:
             if model_class.__name__ == "LevitForImageClassificationWithTeacher":
                  del inputs_dict["labels"]

        return inputs_dict
   def 			_UpperCAmelCase   (			self			)	->      Union[str, Any]:


        '''simple docstring'''


        lowercase_	:				Any													=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__UpperCamelCase			)
   def 			_UpperCAmelCase   (			self			)	->      Optional[int]:


        '''simple docstring'''


        lowercase_	:				Dict													=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*__UpperCamelCase			)
   def 			_UpperCAmelCase   (			self			)	->      Union[str, Any]:


        '''simple docstring'''


        if not self.model_tester.is_training:
             return

        lowercase_							,  lowercase_	:				Union[str, Any]													=	self.model_tester.prepare_config_and_inputs_for_common()
        lowercase_	:				Optional[Any]													=	True

        for model_class in self.all_model_classes:
             # LevitForImageClassificationWithTeacher supports inference-only
             if (
                 model_class in get_values(__UpperCamelCase			)
                 or model_class.__name__ == "LevitForImageClassificationWithTeacher"
             ):
                  continue
             lowercase_	:				Dict													=	model_class(__UpperCamelCase			)
             model.to(__UpperCamelCase			)
             model.train()
             lowercase_	:				List[Any]													=	self._prepare_for_class(__UpperCamelCase     ,__UpperCamelCase     ,return_labels=__UpperCamelCase			)
             lowercase_	:				Union[str, Any]													=	model(**__UpperCamelCase			).loss
             loss.backward()
   def 			_UpperCAmelCase   (			self			)	->      str:


        '''simple docstring'''


        lowercase_							,  lowercase_	:				int													=	self.model_tester.prepare_config_and_inputs_for_common()
        if not self.model_tester.is_training:
             return

        lowercase_	:				Union[str, Any]													=	False
        lowercase_	:				Dict													=	True

        for model_class in self.all_model_classes:
             if model_class in get_values(__UpperCamelCase			) or not model_class.supports_gradient_checkpointing:
                  continue
             # LevitForImageClassificationWithTeacher supports inference-only
             if model_class.__name__ == "LevitForImageClassificationWithTeacher":
                  continue
             lowercase_	:				Dict													=	model_class(__UpperCamelCase			)
             model.gradient_checkpointing_enable()
             model.to(__UpperCamelCase			)
             model.train()
             lowercase_	:				Tuple													=	self._prepare_for_class(__UpperCamelCase     ,__UpperCamelCase     ,return_labels=__UpperCamelCase			)
             lowercase_	:				Any													=	model(**__UpperCamelCase			).loss
             loss.backward()
   def 			_UpperCAmelCase   (			self			)	->      Optional[Any]:


        '''simple docstring'''


        lowercase_							,  lowercase_	:				Any													=	self.model_tester.prepare_config_and_inputs_for_common()

        lowercase_	:				int													=	[
            {'title': 'multi_label_classification', 'num_labels': 2, 'dtype': torch.float},
            {'title': 'single_label_classification', 'num_labels': 1, 'dtype': torch.long},
            {'title': 'regression', 'num_labels': 1, 'dtype': torch.float},
        ]

        for model_class in self.all_model_classes:
             if (
                 model_class
                 not in [
                     *get_values(__UpperCamelCase			),
                 ]
                 or model_class.__name__ == "LevitForImageClassificationWithTeacher"
             ):
                  continue

             for problem_type in problem_types:
                  with self.subTest(msg=f'''Testing {model_class} with {problem_type["title"]}'''			):
                       lowercase_	:				Any													=	problem_type['title']
                       lowercase_	:				Tuple													=	problem_type['num_labels']

                       lowercase_	:				Tuple													=	model_class(__UpperCamelCase			)
                       model.to(__UpperCamelCase			)
                       model.train()

                       lowercase_	:				Optional[int]													=	self._prepare_for_class(__UpperCamelCase     ,__UpperCamelCase     ,return_labels=__UpperCamelCase			)

                       if problem_type["num_labels"] > 1:
                            lowercase_	:				Any													=	inputs['labels'].unsqueeze(1			).repeat(1     ,problem_type['num_labels']			)

                       lowercase_	:				str													=	inputs['labels'].to(problem_type['dtype']			)

                       # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
                       # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
                       # they have the same size." which is a symptom something in wrong for the regression problem.
                       # See https://github.com/huggingface/transformers/issues/11780
                       with warnings.catch_warnings(record=__UpperCamelCase			) as warning_list:
                            lowercase_	:				int													=	model(**__UpperCamelCase			).loss
                       for w in warning_list:
                            if "Using a target size that is different to the input size" in str(w.message			):
                                 raise ValueError(
                                     f'''Something is going wrong in the regression problem: intercepted {w.message}'''			)

                       loss.backward()







   @slow
   def 			_UpperCAmelCase   (			self			)	->      List[Any]:


        '''simple docstring'''


        for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowercase_	:				List[Any]													=	LevitModel.from_pretrained(__UpperCamelCase			)
             self.assertIsNotNone(__UpperCamelCase			)






def 			lowercase__(       ):
     lowercase_	:				Any													=	Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      )
     return image







@require_torch
@require_vision
class        UpperCamelCase		(			unittest.TestCase  ):
   @cached_property
   def 			_UpperCAmelCase   (			self			)	->      Tuple:


        '''simple docstring'''


        return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]			)







   @slow
   def 			_UpperCAmelCase   (			self			)	->      Optional[Any]:


        '''simple docstring'''


        lowercase_	:				Any													=	LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]			).to(
            __UpperCamelCase			)

        lowercase_	:				str													=	self.default_image_processor
        lowercase_	:				Any													=	prepare_img()
        lowercase_	:				Optional[int]													=	image_processor(images=__UpperCamelCase     ,return_tensors='pt'			).to(__UpperCamelCase			)

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

        # verify the logits
        lowercase_	:				str													=	torch.Size((1, 1000)			)
        self.assertEqual(outputs.logits.shape     ,__UpperCamelCase			)

        lowercase_	:				Any													=	torch.tensor([1.0448, -0.3745, -1.8317]			).to(__UpperCamelCase			)

        self.assertTrue(torch.allclose(outputs.logits[0, :3]     ,__UpperCamelCase     ,atol=1e-4			)			)

