

"""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     a  (				lowerCamelCase__ ):



    """simple docstring"""



    UpperCAmelCase         =    42



class     a  (				nn.Module ):



    """simple docstring"""



    def __init__(   self:       str			,  UpperCamelCase:       Tuple=3			,  UpperCamelCase:       List[str]=3			,  UpperCamelCase:       str=("DownEncoderBlock2D",)			,  UpperCamelCase:       List[str]=(64,)			,  UpperCamelCase:       List[str]=2			,  UpperCamelCase:       Optional[Any]=32			,  UpperCamelCase:       List[str]="silu"			,  UpperCamelCase:       List[str]=True			,  ):






         """simple docstring"""


         super().__init__()
         A__								=     layers_per_block

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

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

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

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

         # mid
         A__								=     UNetMidBlockaD(
             in_channels=block_out_channels[-1]			,  resnet_eps=1e-6			,  resnet_act_fn=lowercase__			,  output_scale_factor=1			,  resnet_time_scale_shift="""default"""			,  attention_head_dim=block_out_channels[-1]			,  resnet_groups=lowercase__			,  temb_channels=lowercase__			,  )

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

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

         A__								=     False





    def  UpperCamelCase		(   self:       Any			,  UpperCamelCase:       Any				):






         """simple docstring"""


         A__								=     x
         A__								=     self.conv_in(lowercase__				)

         if self.training and self.gradient_checkpointing:

              def create_custom_forward(UpperCamelCase:       List[Any]				):
                   def custom_forward(*UpperCamelCase:       Dict				):
                        return module(*lowercase__				)

                   return custom_forward

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

         else:
              # down
              for down_block in self.down_blocks:
                   A__								=     down_block(lowercase__				)

              # middle
              A__								=     self.mid_block(lowercase__				)

         # post-process
         A__								=     self.conv_norm_out(lowercase__				)
         A__								=     self.conv_act(lowercase__				)
         A__								=     self.conv_out(lowercase__				)

         return sample



class     a  (				nn.Module ):



    """simple docstring"""



    def __init__(   self:       str			,  UpperCamelCase:       Optional[Any]=3			,  UpperCamelCase:       Any=3			,  UpperCamelCase:       Union[str, Any]=("UpDecoderBlock2D",)			,  UpperCamelCase:       Any=(64,)			,  UpperCamelCase:       str=2			,  UpperCamelCase:       Optional[int]=32			,  UpperCamelCase:       Tuple="silu"			,  UpperCamelCase:       Union[str, Any]="group"			,  ):






         """simple docstring"""


         super().__init__()
         A__								=     layers_per_block

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

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

         A__								=     in_channels if norm_type == """spatial""" else None

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

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

              A__								=     i == len(lowercase__				) - 1

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

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

         A__								=     False





    def  UpperCamelCase		(   self:       str			,  UpperCamelCase:       Optional[Any]			,  UpperCamelCase:       List[Any]=None				):






         """simple docstring"""


         A__								=     z
         A__								=     self.conv_in(lowercase__				)

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

              def create_custom_forward(UpperCamelCase:       str				):
                   def custom_forward(*UpperCamelCase:       List[Any]				):
                        return module(*lowercase__				)

                   return custom_forward

              if is_torch_version(""">="""			,  """1.11.0"""				):
                   # middle
                   A__								=     torch.utils.checkpoint.checkpoint(
                       create_custom_forward(self.mid_block				)			,  lowercase__			,  lowercase__			,  use_reentrant=lowercase__				)
                   A__								=     sample.to(lowercase__				)

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

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

              # up
              for up_block in self.up_blocks:
                   A__								=     up_block(lowercase__			,  lowercase__				)

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

         return sample



class     a  (				nn.Module ):



    """simple docstring"""



    def __init__(   self:       int			,  UpperCamelCase:       str			,  UpperCamelCase:       List[Any]			,  UpperCamelCase:       str			,  UpperCamelCase:       int=None			,  UpperCamelCase:       int="random"			,  UpperCamelCase:       Tuple=False			,  UpperCamelCase:       List[str]=True				):






         """simple docstring"""


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

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

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

         A__								=     sane_index_shape

    def  UpperCamelCase		(   self:       List[str]			,  UpperCamelCase:       List[Any]				):






         """simple docstring"""


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

    def  UpperCamelCase		(   self:       Dict			,  UpperCamelCase:       Any				):






         """simple docstring"""


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

    def  UpperCamelCase		(   self:       Optional[Any]			,  UpperCamelCase:       Tuple				):






         """simple docstring"""


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

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

         A__								=     self.embedding(lowercase__				).view(z.shape				)
         A__								=     None
         A__								=     None

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

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

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

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

         if self.sane_index_shape:
              A__								=     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  UpperCamelCase		(   self:       Dict			,  UpperCamelCase:       List[Any]			,  UpperCamelCase:       List[str]				):






         """simple docstring"""


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

         # get quantized latent vectors
         A__								=     self.embedding(lowercase__				)

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

         return z_q



class     a  (				lowerCamelCase__ ):



    """simple docstring"""



    def __init__(   self:       str			,  UpperCamelCase:       Union[str, Any]			,  UpperCamelCase:       Union[str, Any]=False				):






         """simple docstring"""


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

    def  UpperCamelCase		(   self:       Union[str, Any]			,  UpperCamelCase:       Optional[torch.Generator] = None				):






         """simple docstring"""


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

    def  UpperCamelCase		(   self:       Dict			,  UpperCamelCase:       Any=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  UpperCamelCase		(   self:       Tuple			,  UpperCamelCase:       List[str]			,  UpperCamelCase:       str=[1, 2, 3]				):






         """simple docstring"""


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





    def  UpperCamelCase		(   self:       str				):






         """simple docstring"""


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




  """simple docstring"""




  def __init__(      self   :     str					,lowercase__   :     Tuple					,lowercase__   :     Dict=1_3					,lowercase__   :     List[str]=3_0					,lowercase__   :     Tuple=2					,lowercase__   :     Optional[int]=3					,lowercase__   :     List[str]=True					,lowercase__   :     Tuple=True					,lowercase__   :     int=3_2					,lowercase__   :     List[str]=5					,lowercase__   :     Tuple=4					,lowercase__   :     Any=3_7					,lowercase__   :     Any="gelu"					,lowercase__   :     Union[str, Any]=0.1					,lowercase__   :     Optional[int]=0.1					,lowercase__   :     str=1_0					,lowercase__   :     Optional[int]=0.0_2					,lowercase__   :     Union[str, Any]=3					,lowercase__   :     Optional[int]=0.6					,lowercase__   :     List[Any]=None					,):
   __lowercase							  = parent
   __lowercase							  = batch_size
   __lowercase							  = image_size
   __lowercase							  = patch_size
   __lowercase							  = num_channels
   __lowercase							  = is_training
   __lowercase							  = use_labels
   __lowercase							  = hidden_size
   __lowercase							  = num_hidden_layers
   __lowercase							  = num_attention_heads
   __lowercase							  = intermediate_size
   __lowercase							  = hidden_act
   __lowercase							  = hidden_dropout_prob
   __lowercase							  = attention_probs_dropout_prob
   __lowercase							  = type_sequence_label_size
   __lowercase							  = initializer_range
   __lowercase							  = mask_ratio
   __lowercase							  = scope

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






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]			):
   __lowercase							  = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

   __lowercase							  = None
   if self.use_labels:
    __lowercase							  = ids_tensor([self.batch_size]					,self.type_sequence_label_size			)

   __lowercase							  = self.get_config()

   return config, pixel_values, labels






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]			):
   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=lowercase__					,initializer_range=self.initializer_range					,mask_ratio=self.mask_ratio					,)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[Any]					,lowercase__   :     str					,lowercase__   :     Optional[int]					,lowercase__   :     List[str]			):
   __lowercase							  = ViTMAEModel(config=lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__			)
   self.parent.assertEqual(result.last_hidden_state.shape					,(self.batch_size, self.seq_length, self.hidden_size)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Any					,lowercase__   :     List[Any]					,lowercase__   :     int					,lowercase__   :     Optional[Any]			):
   __lowercase							  = ViTMAEForPreTraining(lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__			)
   __lowercase							  = (self.image_size // self.patch_size) ** 2
   __lowercase							  = self.patch_size**2 * self.num_channels
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, num_patches, expected_num_channels)			)

   # test greyscale images
   __lowercase							  = 1
   __lowercase							  = ViTMAEForPreTraining(lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
   __lowercase							  = model(lowercase__			)
   __lowercase							  = self.patch_size**2
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, num_patches, expected_num_channels)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]			):
   __lowercase							  = self.prepare_config_and_inputs()
   __lowercase  ,	__lowercase  ,	__lowercase							  = config_and_inputs
   __lowercase							  = {'''pixel_values''': pixel_values}
   return config, inputs_dict






@require_torch
class 			lowercase_	(lowerCamelCase__				,       lowerCamelCase__				,       unittest.TestCase      ):




  """simple docstring"""




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

  SCREAMING_SNAKE_CASE    :  Any							=			False
  SCREAMING_SNAKE_CASE    :  Optional[int]							=			False
  SCREAMING_SNAKE_CASE    :  List[str]							=			False
  SCREAMING_SNAKE_CASE    :  Union[str, Any]							=			False






  def 					SCREAMING_SNAKE_CASE      (      self   :     int			):
   __lowercase							  = ViTMAEModelTester(self			)
   __lowercase							  = ConfigTester(self					,config_class=lowercase__					,has_text_modality=lowercase__					,hidden_size=3_7			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]			):
   self.config_tester.run_common_tests()






  @unittest.skip(reason='''ViTMAE does not use inputs_embeds'''			)
  def 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   pass






  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   __lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs_for_common()

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






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]			):
   __lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs_for_common()

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

    __lowercase							  = ['''pixel_values''']
    self.assertListEqual(arg_names[:1]					,lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[Any]			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_pretraining(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Any					,lowercase__   :     Optional[int]					,lowercase__   :     List[str]					,lowercase__   :     Optional[Any]			):
   # make masks reproducible
   np.random.seed(2			)

   __lowercase							  = int((pt_model.config.image_size // pt_model.config.patch_size) ** 2			)
   __lowercase							  = np.random.uniform(size=(self.model_tester.batch_size, num_patches)			)
   __lowercase							  = torch.from_numpy(lowercase__			)

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

   super().check_pt_tf_models(lowercase__					,lowercase__					,lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]			):
   __lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    __lowercase							  = model_class(lowercase__			)
    model.to(lowercase__			)
    model.eval()
    # make random mask reproducible
    torch.manual_seed(2			)
    with torch.no_grad():
     __lowercase							  = model(**self._prepare_for_class(lowercase__					,lowercase__			)			)

    __lowercase							  = outputs[0].cpu().numpy()
    __lowercase							  = 0

    with tempfile.TemporaryDirectory() as tmpdirname:
     model.save_pretrained(lowercase__			)
     __lowercase							  = model_class.from_pretrained(lowercase__			)
     model.to(lowercase__			)
     # make random mask reproducible
     torch.manual_seed(2			)
     with torch.no_grad():
      __lowercase							  = model(**self._prepare_for_class(lowercase__					,lowercase__			)			)

     # Make sure we don't have nans
     __lowercase							  = after_outputs[0].cpu().numpy()
     __lowercase							  = 0
     __lowercase							  = np.amax(np.abs(out_a - out_a			)			)
     self.assertLessEqual(lowercase__					,1e-5			)






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






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






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






  @unittest.skip(reason='''ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load'''			)
  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]			):
   pass






  @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''			)
  def 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   pass






  @slow
  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]			):
   for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    __lowercase							  = ViTMAEModel.from_pretrained(lowercase__			)
    self.assertIsNotNone(lowercase__			)


def 		_A     (    ):






 """simple docstring"""



 __lowercase							  = Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''      )
 return image






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




  """simple docstring"""




  @cached_property
  def 					SCREAMING_SNAKE_CASE      (      self   :     str			):
   return ViTImageProcessor.from_pretrained('''facebook/vit-mae-base'''			) if is_vision_available() else None






  @slow
  def 					SCREAMING_SNAKE_CASE      (      self   :     str			):
   # make random mask reproducible across the PT and TF model
   np.random.seed(2			)

   __lowercase							  = ViTMAEForPreTraining.from_pretrained('''facebook/vit-mae-base'''			).to(lowercase__			)

   __lowercase							  = self.default_image_processor
   __lowercase							  = prepare_img()
   __lowercase							  = image_processor(images=lowercase__					,return_tensors='''pt'''			).to(lowercase__			)

   # 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)
   __lowercase							  = ViTMAEConfig()
   __lowercase							  = int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2			)
   __lowercase							  = np.random.uniform(size=(1, num_patches)			)

   # forward pass
   with torch.no_grad():
    __lowercase							  = model(**lowercase__					,noise=torch.from_numpy(lowercase__			).to(device=lowercase__			)			)

   # verify the logits
   __lowercase							  = torch.Size((1, 1_9_6, 7_6_8)			)
   self.assertEqual(outputs.logits.shape					,lowercase__			)

   __lowercase							  = torch.tensor(
       [[-0.0_5_4_8, -1.7_0_2_3, -0.9_3_2_5], [0.3_7_2_1, -0.5_6_7_0, -0.2_2_3_3], [0.8_2_3_5, -1.3_8_7_8, -0.3_5_2_4]]			)

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

