





'''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     lowerCAmelCase__ (				UpperCAmelCase_    ):

       '''simple docstring'''
       _lowerCamelCase								=42






class     lowerCAmelCase__ (				nn.Module    ):

       '''simple docstring'''

       def __init__( self					:				Tuple       ,     a__					:				str=3       ,     a__					:				Optional[int]=3       ,     a__					:				str=("DownEncoderBlock2D",)       ,     a__					:				Optional[int]=(64,)       ,     a__					:				List[Any]=2       ,     a__					:				Tuple=32       ,     a__					:				List[Any]="silu"       ,     a__					:				Dict=True       ,     ):
              super().__init__()
              UpperCAmelCase  =						layers_per_block

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

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

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

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

              # mid
              UpperCAmelCase  =						UNetMidBlockaD(
                  in_channels=block_out_channels[-1]       ,     resnet_eps=1e-6       ,     resnet_act_fn=a__       ,     output_scale_factor=1       ,     resnet_time_scale_shift='''default'''       ,     attention_head_dim=block_out_channels[-1]       ,     resnet_groups=a__       ,     temb_channels=a__       ,     )

              # out
              UpperCAmelCase  =						nn.GroupNorm(num_channels=block_out_channels[-1]       ,     num_groups=a__       ,     eps=1e-6						)
              UpperCAmelCase  =						nn.SiLU()

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

              UpperCAmelCase  =						False



       def 			__snake_case      ( self					:				Any       ,     a__					:				Optional[int]						):
              UpperCAmelCase  =						x
              UpperCAmelCase  =						self.conv_in(a__						)

              if self.training and self.gradient_checkpointing:

                     def create_custom_forward(a__					:				Dict						):
                            def custom_forward(*a__					:				Union[str, Any]						):
                                   return module(*a__						)

                            return custom_forward

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

              else:
                     # down
                     for down_block in self.down_blocks:
                            UpperCAmelCase  =						down_block(a__						)

                     # middle
                     UpperCAmelCase  =						self.mid_block(a__						)

              # post-process
              UpperCAmelCase  =						self.conv_norm_out(a__						)
              UpperCAmelCase  =						self.conv_act(a__						)
              UpperCAmelCase  =						self.conv_out(a__						)

              return sample






class     lowerCAmelCase__ (				nn.Module    ):

       '''simple docstring'''

       def __init__( self					:				List[str]       ,     a__					:				Optional[int]=3       ,     a__					:				List[str]=3       ,     a__					:				Dict=("UpDecoderBlock2D",)       ,     a__					:				List[str]=(64,)       ,     a__					:				Union[str, Any]=2       ,     a__					:				Optional[Any]=32       ,     a__					:				Union[str, Any]="silu"       ,     a__					:				str="group"       ,     ):
              super().__init__()
              UpperCAmelCase  =						layers_per_block

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

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

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

              # mid
              UpperCAmelCase  =						UNetMidBlockaD(
                  in_channels=block_out_channels[-1]       ,     resnet_eps=1e-6       ,     resnet_act_fn=a__       ,     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=a__       ,     temb_channels=a__       ,     )

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

                     UpperCAmelCase  =						i == len(a__						) - 1

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

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

              UpperCAmelCase  =						False



       def 			__snake_case      ( self					:				Dict       ,     a__					:				Optional[Any]       ,     a__					:				Any=None						):
              UpperCAmelCase  =						z
              UpperCAmelCase  =						self.conv_in(a__						)

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

                     def create_custom_forward(a__					:				Optional[int]						):
                            def custom_forward(*a__					:				Tuple						):
                                   return module(*a__						)

                            return custom_forward

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

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

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

                     # up
                     for up_block in self.up_blocks:
                            UpperCAmelCase  =						up_block(a__       ,     a__						)

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

              return sample






class     lowerCAmelCase__ (				nn.Module    ):

       '''simple docstring'''

       def __init__( self					:				List[Any]       ,     a__					:				Optional[int]       ,     a__					:				Optional[int]       ,     a__					:				str       ,     a__					:				List[str]=None       ,     a__					:				str="random"       ,     a__					:				List[str]=False       ,     a__					:				Tuple=True						):
              super().__init__()
              UpperCAmelCase  =						n_e
              UpperCAmelCase  =						vq_embed_dim
              UpperCAmelCase  =						beta
              UpperCAmelCase  =						legacy

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

              UpperCAmelCase  =						remap
              if self.remap is not None:
                     self.register_buffer('''used'''       ,     torch.tensor(np.load(self.remap						)						)						)
                     UpperCAmelCase  =						self.used.shape[0]
                     UpperCAmelCase  =						unknown_index  # "random" or "extra" or integer
                     if self.unknown_index == "extra":
                            UpperCAmelCase  =						self.re_embed
                            UpperCAmelCase  =						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:
                     UpperCAmelCase  =						n_e

              UpperCAmelCase  =						sane_index_shape

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

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

       def 			__snake_case      ( self					:				List[Any]       ,     a__					:				Optional[Any]						):
              # reshape z -> (batch, height, width, channel) and flatten
              UpperCAmelCase  =						z.permute(0       ,     2       ,     3       ,     1						).contiguous()
              UpperCAmelCase  =						z.view(-1       ,     self.vq_embed_dim						)

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

              UpperCAmelCase  =						self.embedding(a__						).view(z.shape						)
              UpperCAmelCase  =						None
              UpperCAmelCase  =						None

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

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

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

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

              if self.sane_index_shape:
                     UpperCAmelCase  =						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 			__snake_case      ( self					:				str       ,     a__					:				Any       ,     a__					:				str						):
              # shape specifying (batch, height, width, channel)
              if self.remap is not None:
                     UpperCAmelCase  =						indices.reshape(shape[0]       ,     -1						)  # add batch axis
                     UpperCAmelCase  =						self.unmap_to_all(a__						)
                     UpperCAmelCase  =						indices.reshape(-1						)  # flatten again

              # get quantized latent vectors
              UpperCAmelCase  =						self.embedding(a__						)

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

              return z_q






class     lowerCAmelCase__ (				UpperCAmelCase_    ):

       '''simple docstring'''

       def __init__( self					:				Optional[Any]       ,     a__					:				Union[str, Any]       ,     a__					:				Union[str, Any]=False						):
              UpperCAmelCase  =						parameters
              UpperCAmelCase,					UpperCAmelCase  =						torch.chunk(a__       ,     2       ,     dim=1						)
              UpperCAmelCase  =						torch.clamp(self.logvar       ,     -30.0       ,     20.0						)
              UpperCAmelCase  =						deterministic
              UpperCAmelCase  =						torch.exp(0.5 * self.logvar						)
              UpperCAmelCase  =						torch.exp(self.logvar						)
              if self.deterministic:
                     UpperCAmelCase  =						UpperCAmelCase  =						torch.zeros_like(
                         self.mean       ,     device=self.parameters.device       ,     dtype=self.parameters.dtype						)

       def 			__snake_case      ( self					:				Optional[Any]       ,     a__					:				Optional[torch.Generator] = None						):
              # make sure sample is on the same device as the parameters and has same dtype
              UpperCAmelCase  =						randn_tensor(
                  self.mean.shape       ,     generator=a__       ,     device=self.parameters.device       ,     dtype=self.parameters.dtype						)
              UpperCAmelCase  =						self.mean + self.std * sample
              return x

       def 			__snake_case      ( self					:				List[Any]       ,     a__					:				str=None						):
              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 			__snake_case      ( self					:				Tuple       ,     a__					:				str       ,     a__					:				List[Any]=[1, 2, 3]						):
              if self.deterministic:
                     return torch.Tensor([0.0]						)
              UpperCAmelCase  =						np.log(2.0 * np.pi						)
              return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean       ,     2						) / self.var       ,     dim=a__						)



       def 			__snake_case      ( self					:				Optional[int]						):
              return self.mean






'''simple docstring'''


import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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 torch import nn

      from transformers import ViTMAEForPreTraining, ViTMAEModel
      from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import ViTImageProcessor







class 							__magic_name__	:

     """simple docstring"""






     def __init__(   self			,     lowerCamelCase			,     lowerCamelCase=13			,     lowerCamelCase=30			,     lowerCamelCase=2			,     lowerCamelCase=3			,     lowerCamelCase=True			,     lowerCamelCase=True			,     lowerCamelCase=32			,     lowerCamelCase=5			,     lowerCamelCase=4			,     lowerCamelCase=37			,     lowerCamelCase="gelu"			,     lowerCamelCase=0.1			,     lowerCamelCase=0.1			,     lowerCamelCase=10			,     lowerCamelCase=0.02			,     lowerCamelCase=3			,     lowerCamelCase=0.6			,     lowerCamelCase=None			,     ):


          '''simple docstring'''





          __A	:       Tuple      =							parent
          __A	:       Union[str, Any]      =							batch_size
          __A	:       List[Any]      =							image_size
          __A	:       Union[str, Any]      =							patch_size
          __A	:       List[Any]      =							num_channels
          __A	:       Optional[Any]      =							is_training
          __A	:       str      =							use_labels
          __A	:       Tuple      =							hidden_size
          __A	:       int      =							num_hidden_layers
          __A	:       Dict      =							num_attention_heads
          __A	:       List[Any]      =							intermediate_size
          __A	:       Tuple      =							hidden_act
          __A	:       Tuple      =							hidden_dropout_prob
          __A	:       str      =							attention_probs_dropout_prob
          __A	:       Optional[Any]      =							type_sequence_label_size
          __A	:       Union[str, Any]      =							initializer_range
          __A	:       Optional[Any]      =							mask_ratio
          __A	:       List[Any]      =							scope

          # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
          # (we add 1 for the [CLS] token)
          __A	:       Optional[int]      =							(image_size // patch_size) ** 2
          __A	:       List[Any]      =							int(math.ceil((1 - mask_ratio) * (num_patches + 1)		)		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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

          __A	:       Tuple      =							None
          if self.use_labels:
               __A	:       Union[str, Any]      =							ids_tensor([self.batch_size]			,     self.type_sequence_label_size		)

          __A	:       Optional[int]      =							self.get_config()

          return config, pixel_values, labels


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          return ViTMAEConfig(
              image_size=self.image_size			,     patch_size=self.patch_size			,     num_channels=self.num_channels			,     hidden_size=self.hidden_size			,     num_hidden_layers=self.num_hidden_layers			,     num_attention_heads=self.num_attention_heads			,     intermediate_size=self.intermediate_size			,     hidden_act=self.hidden_act			,     hidden_dropout_prob=self.hidden_dropout_prob			,     attention_probs_dropout_prob=self.attention_probs_dropout_prob			,     is_decoder=lowerCamelCase			,     initializer_range=self.initializer_range			,     mask_ratio=self.mask_ratio			,     )


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       Optional[Any]      =							ViTMAEModel(config=lowerCamelCase		)
          model.to(lowerCamelCase		)
          model.eval()
          __A	:       List[Any]      =							model(lowerCamelCase		)
          self.parent.assertEqual(result.last_hidden_state.shape			,     (self.batch_size, self.seq_length, self.hidden_size)		)


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		):


          '''simple docstring'''





          __A	:       str      =							ViTMAEForPreTraining(lowerCamelCase		)
          model.to(lowerCamelCase		)
          model.eval()
          __A	:       Optional[Any]      =							model(lowerCamelCase		)
          __A	:       List[str]      =							(self.image_size // self.patch_size) ** 2
          __A	:       List[Any]      =							self.patch_size**2 * self.num_channels
          self.parent.assertEqual(result.logits.shape			,     (self.batch_size, num_patches, expected_num_channels)		)

          # test greyscale images
          __A	:       List[str]      =							1
          __A	:       str      =							ViTMAEForPreTraining(lowerCamelCase		)
          model.to(lowerCamelCase		)
          model.eval()
          __A	:       str      =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
          __A	:       int      =							model(lowerCamelCase		)
          __A	:       Union[str, Any]      =							self.patch_size**2
          self.parent.assertEqual(result.logits.shape			,     (self.batch_size, num_patches, expected_num_channels)		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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







@require_torch
class 							__magic_name__	(       lowerCAmelCase    ,		lowerCAmelCase    ,		unittest.TestCase ):

     """simple docstring"""






     lowerCamelCase__       						=				(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
     lowerCamelCase__       						=				{'feature-extraction': ViTMAEModel} if is_torch_available() else {}

     lowerCamelCase__       						=				False
     lowerCamelCase__       						=				False
     lowerCamelCase__       						=				False
     lowerCamelCase__       						=				False


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Union[str, Any]      =							ViTMAEModelTester(self		)
          __A	:       Tuple      =							ConfigTester(self			,     config_class=lowerCamelCase			,     has_text_modality=lowerCamelCase			,     hidden_size=37		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          self.config_tester.run_common_tests()


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


          '''simple docstring'''





          pass


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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

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


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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

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

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


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          __A	:       Union[str, Any]      =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_pretraining(*lowerCamelCase		)


     def lowerCAmelCase__  (   self			,     lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		):


          '''simple docstring'''





          np.random.seed(2		)

          __A	:       Tuple      =							int((pt_model.config.image_size // pt_model.config.patch_size) ** 2		)
          __A	:       Optional[Any]      =							np.random.uniform(size=(self.model_tester.batch_size, num_patches)		)
          __A	:       Optional[Any]      =							torch.from_numpy(lowerCamelCase		)

          # Add `noise` argument.
          # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
          __A	:       int      =							pt_noise

          super().check_pt_tf_models(lowerCamelCase			,     lowerCamelCase			,     lowerCamelCase		)


     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





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

          for model_class in self.all_model_classes:
               __A	:       int      =							model_class(lowerCamelCase		)
               model.to(lowerCamelCase		)
               model.eval()
               # make random mask reproducible
               torch.manual_seed(2		)
               with torch.no_grad():
                    __A	:       List[str]      =							model(**self._prepare_for_class(lowerCamelCase			,     lowerCamelCase		)		)

               __A	:       Tuple      =							outputs[0].cpu().numpy()
               __A	:       List[Any]      =							0

               with tempfile.TemporaryDirectory() as tmpdirname:
                    model.save_pretrained(lowerCamelCase		)
                    __A	:       List[Any]      =							model_class.from_pretrained(lowerCamelCase		)
                    model.to(lowerCamelCase		)
                    # make random mask reproducible
                    torch.manual_seed(2		)
                    with torch.no_grad():
                         __A	:       Optional[Any]      =							model(**self._prepare_for_class(lowerCamelCase			,     lowerCamelCase		)		)

                    # Make sure we don't have nans
                    __A	:       List[Any]      =							after_outputs[0].cpu().numpy()
                    __A	:       List[str]      =							0
                    __A	:       int      =							np.amax(np.abs(out_a - out_a		)		)
                    self.assertLessEqual(lowerCamelCase			,     1E-5		)


     @unittest.skip(
         reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."		)
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          pass


     @unittest.skip(
         reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."		)
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          pass


     @unittest.skip(
         reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."		)
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          pass


     @unittest.skip(reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"		)
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          pass


     @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."		)
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          pass


     @slow
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __A	:       List[Any]      =							ViTMAEModel.from_pretrained(lowerCamelCase		)
               self.assertIsNotNone(lowerCamelCase		)


def 		_lowercase ():




     '''simple docstring'''





     __A	:       Union[str, Any]      =							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"							)
     return image







@require_torch
@require_vision
class 							__magic_name__	(       unittest.TestCase ):

     """simple docstring"""






     @cached_property
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          return ViTImageProcessor.from_pretrained("facebook/vit-mae-base"		) if is_vision_available() else None


     @slow
     def lowerCAmelCase__  (   self		):


          '''simple docstring'''





          np.random.seed(2		)

          __A	:       Dict      =							ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"		).to(lowerCamelCase		)

          __A	:       str      =							self.default_image_processor
          __A	:       List[Any]      =							prepare_img()
          __A	:       Union[str, Any]      =							image_processor(images=lowerCamelCase			,     return_tensors="pt"		).to(lowerCamelCase		)

          # prepare a noise vector that will be also used for testing the TF model
          # (this way we can ensure that the PT and TF models operate on the same inputs)
          __A	:       Optional[Any]      =							ViTMAEConfig()
          __A	:       List[str]      =							int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2		)
          __A	:       Any      =							np.random.uniform(size=(1, num_patches)		)

          # forward pass
          with torch.no_grad():
               __A	:       str      =							model(**lowerCamelCase			,     noise=torch.from_numpy(lowerCamelCase		).to(device=lowerCamelCase		)		)

          # verify the logits
          __A	:       Union[str, Any]      =							torch.Size((1, 196, 768)		)
          self.assertEqual(outputs.logits.shape			,     lowerCamelCase		)

          __A	:       int      =							torch.tensor(
              [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]]		)

          self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]			,     expected_slice.to(lowerCamelCase		)			,     atol=1E-4		)		)

