





import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


_lowerCAmelCase 				=     logging.get_logger(__name__)

# General docstring
_lowerCAmelCase 				=     """PoolFormerConfig"""

# Base docstring
_lowerCAmelCase 				=     """sail/poolformer_s12"""
_lowerCAmelCase 				=     [1, 512, 7, 7]

# Image classification docstring
_lowerCAmelCase 				=     """sail/poolformer_s12"""
_lowerCAmelCase 				=     """tabby, tabby cat"""

_lowerCAmelCase 				=     [
    """sail/poolformer_s12""",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]



def     _lowerCAmelCase		(  _lowerCAmelCase  ,_lowerCAmelCase = 0.0  ,_lowerCAmelCase = False     ):
   '''simple docstring'''


   if drop_prob == 0.0 or not training:
      return input
   A_     :				int			      =	1 - drop_prob
   A_     :				List[str]			      =	(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
   A_     :				Optional[int]			      =	keep_prob + torch.rand(_lowerCAmelCase  ,dtype=input.dtype  ,device=input.device     )
   random_tensor.floor_()  # binarize
   A_     :				Optional[int]			      =	input.div(_lowerCAmelCase     ) * random_tensor
   return output

class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__ = None       ):
         super().__init__()
         A_     :				int			      =	drop_prob
      def 		_lowerCamelCase      (		self      ,       a__       ):
         return drop_path(a__      ,       self.drop_prob      ,       self.training       )






      def 		_lowerCamelCase      (		self       ):
         return "p={}".format(self.drop_prob       )






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__=None       ):
         super().__init__()
         A_     :				Dict			      =	patch_size if isinstance(a__      ,       collections.abc.Iterable       ) else (patch_size, patch_size)
         A_     :				Optional[Any]			      =	stride if isinstance(a__      ,       collections.abc.Iterable       ) else (stride, stride)
         A_     :				int			      =	padding if isinstance(a__      ,       collections.abc.Iterable       ) else (padding, padding)

         A_     :				Dict			      =	nn.Convad(a__      ,       a__      ,       kernel_size=a__      ,       stride=a__      ,       padding=a__       )
         A_     :				Optional[Any]			      =	norm_layer(a__       ) if norm_layer else nn.Identity()






      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				Dict			      =	self.projection(a__       )
         A_     :				List[Any]			      =	self.norm(a__       )
         return embeddings






class    _UpperCAmelCase							(		nn.GroupNorm     ):
      def __init__(		self      ,       a__      ,       **a__       ):
         super().__init__(1      ,       a__      ,       **a__       )






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__       ):
         super().__init__()
         A_     :				Optional[int]			      =	nn.AvgPoolad(a__      ,       stride=1      ,       padding=pool_size // 2      ,       count_include_pad=a__       )






      def 		_lowerCamelCase      (		self      ,       a__       ):
         return self.pool(a__       ) - hidden_states






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__      ,       a__      ,       a__      ,       a__       ):
         super().__init__()
         A_     :				str			      =	nn.Convad(a__      ,       a__      ,       1       )
         A_     :				List[str]			      =	nn.Convad(a__      ,       a__      ,       1       )
         A_     :				Optional[Any]			      =	PoolFormerDropPath(a__       )
         if isinstance(config.hidden_act      ,       a__       ):
            A_     :				int			      =	ACTaFN[config.hidden_act]
         else:
            A_     :				Optional[int]			      =	config.hidden_act






      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				Tuple			      =	self.conva(a__       )
         A_     :				str			      =	self.act_fn(a__       )
         A_     :				Union[str, Any]			      =	self.drop(a__       )
         A_     :				int			      =	self.conva(a__       )
         A_     :				Any			      =	self.drop(a__       )

         return hidden_states






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__       ):
         super().__init__()
         A_     :				Dict			      =	PoolFormerPooling(a__       )
         A_     :				Dict			      =	PoolFormerOutput(a__      ,       a__      ,       a__      ,       a__       )
         A_     :				Dict			      =	PoolFormerGroupNorm(a__       )
         A_     :				Dict			      =	PoolFormerGroupNorm(a__       )

         # Useful for training neural nets
         A_     :				Union[str, Any]			      =	PoolFormerDropPath(a__       ) if drop_path > 0.0 else nn.Identity()
         A_     :				Union[str, Any]			      =	config.use_layer_scale
         if config.use_layer_scale:
            A_     :				Tuple			      =	nn.Parameter(
                config.layer_scale_init_value * torch.ones((a__)       )      ,       requires_grad=a__       )
            A_     :				Tuple			      =	nn.Parameter(
                config.layer_scale_init_value * torch.ones((a__)       )      ,       requires_grad=a__       )






      def 		_lowerCamelCase      (		self      ,       a__       ):
         if self.use_layer_scale:
            A_     :				List[str]			      =	self.pooling(self.before_norm(a__       )       )
            A_     :				Optional[Any]			      =	self.layer_scale_a.unsqueeze(-1       ).unsqueeze(-1       ) * pooling_output
            # First residual connection
            A_     :				List[Any]			      =	hidden_states + self.drop_path(a__       )
            A_     :				Tuple			      =	()

            A_     :				List[Any]			      =	self.output(self.after_norm(a__       )       )
            A_     :				str			      =	self.layer_scale_a.unsqueeze(-1       ).unsqueeze(-1       ) * layer_output
            # Second residual connection
            A_     :				Optional[Any]			      =	hidden_states + self.drop_path(a__       )

            A_     :				int			      =	(output,) + outputs
            return outputs

         else:
            A_     :				Union[str, Any]			      =	self.drop_path(self.pooling(self.before_norm(a__       )       )       )
            # First residual connection
            A_     :				Any			      =	pooling_output + hidden_states
            A_     :				Tuple			      =	()

            # Second residual connection inside the PoolFormerOutput block
            A_     :				Dict			      =	self.drop_path(self.output(self.after_norm(a__       )       )       )
            A_     :				int			      =	hidden_states + layer_output

            A_     :				str			      =	(output,) + outputs
            return outputs






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__       ):
         super().__init__()
         A_     :				Union[str, Any]			      =	config
         # stochastic depth decay rule
         A_     :				Union[str, Any]			      =	[x.item() for x in torch.linspace(0      ,       config.drop_path_rate      ,       sum(config.depths       )       )]

         # patch embeddings
         A_     :				Union[str, Any]			      =	[]
         for i in range(config.num_encoder_blocks       ):
            embeddings.append(
                PoolFormerEmbeddings(
                    patch_size=config.patch_sizes[i]      ,       stride=config.strides[i]      ,       padding=config.padding[i]      ,       num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]      ,       hidden_size=config.hidden_sizes[i]      ,       )       )
         A_     :				Optional[int]			      =	nn.ModuleList(a__       )

         # Transformer blocks
         A_     :				Union[str, Any]			      =	[]
         A_     :				Union[str, Any]			      =	0
         for i in range(config.num_encoder_blocks       ):
            # each block consists of layers
            A_     :				Dict			      =	[]
            if i != 0:
               cur += config.depths[i - 1]
            for j in range(config.depths[i]       ):
               layers.append(
                   PoolFormerLayer(
                       a__      ,       num_channels=config.hidden_sizes[i]      ,       pool_size=config.pool_size      ,       hidden_size=config.hidden_sizes[i]      ,       intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio       )      ,       drop_path=dpr[cur + j]      ,       )       )
            blocks.append(nn.ModuleList(a__       )       )

         A_     :				str			      =	nn.ModuleList(a__       )






      def 		_lowerCamelCase      (		self      ,       a__      ,       a__=False      ,       a__=True       ):
         A_     :				Optional[int]			      =	() if output_hidden_states else None

         A_     :				str			      =	pixel_values
         for idx, layers in enumerate(zip(self.patch_embeddings      ,       self.block       )       ):
            A_      ,				A_     :				Union[str, Any]			      =	layers
            # Get patch embeddings from hidden_states
            A_     :				Tuple			      =	embedding_layer(a__       )
            # Send the embeddings through the blocks
            for _, blk in enumerate(a__       ):
               A_     :				List[str]			      =	blk(a__       )
               A_     :				Union[str, Any]			      =	layer_outputs[0]

            if output_hidden_states:
               A_     :				Dict			      =	all_hidden_states + (hidden_states,)

         if not return_dict:
            return tuple(v for v in [hidden_states, all_hidden_states] if v is not None       )

         return BaseModelOutputWithNoAttention(last_hidden_state=a__      ,       hidden_states=a__       )






class    _UpperCAmelCase							(		_lowerCamelCase     ):
      a						       =							PoolFormerConfig
      a						       =							'''poolformer'''
      a						       =							'''pixel_values'''
      a						       =							True
      def 		_lowerCamelCase      (		self      ,       a__       ):
         if isinstance(a__      ,       (nn.Linear, nn.Convad)       ):
            module.weight.data.normal_(mean=0.0      ,       std=self.config.initializer_range       )
            if module.bias is not None:
               module.bias.data.zero_()
         elif isinstance(a__      ,       nn.LayerNorm       ):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0       )






      def 		_lowerCamelCase      (		self      ,       a__      ,       a__=False       ):
         if isinstance(a__      ,       a__       ):
            A_     :				List[Any]			      =	value






_lowerCAmelCase 				=     R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

_lowerCAmelCase 				=     R"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`PoolFormerImageProcessor.__call__`] for details.
"""

@add_start_docstrings(
    '''The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top.'''			,       _lowerCamelCase			,       )
class    _UpperCAmelCase							(		_lowerCamelCase     ):
      def __init__(		self      ,       a__       ):
         super().__init__(a__       )
         A_     :				List[Any]			      =	config

         A_     :				Optional[int]			      =	PoolFormerEncoder(a__       )

         # Initialize weights and apply final processing
         self.post_init()
      def 		_lowerCamelCase      (		self       ):
         return self.embeddings.patch_embeddings






      @add_start_docstrings_to_model_forward(a__       )
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC      ,       output_type=a__      ,       config_class=_CONFIG_FOR_DOC      ,       modality="""vision"""      ,       expected_output=_EXPECTED_OUTPUT_SHAPE      ,       )
      def 		_lowerCamelCase      (		self      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       ):
         A_     :				Union[str, Any]			      =	(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         A_     :				str			      =	return_dict if return_dict is not None else self.config.use_return_dict

         if pixel_values is None:
            raise ValueError("""You have to specify pixel_values"""       )

         A_     :				Optional[Any]			      =	self.encoder(
             a__      ,       output_hidden_states=a__      ,       return_dict=a__      ,       )
         A_     :				Optional[int]			      =	encoder_outputs[0]

         if not return_dict:
            return (sequence_output, None) + encoder_outputs[1:]

         return BaseModelOutputWithNoAttention(
             last_hidden_state=a__      ,       hidden_states=encoder_outputs.hidden_states      ,       )






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__       ):
         super().__init__()
         A_     :				int			      =	nn.Linear(config.hidden_size      ,       config.hidden_size       )






      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				Any			      =	self.dense(a__       )
         return output






@add_start_docstrings(
    '''
    PoolFormer Model transformer with an image classification head on top
    '''			,       _lowerCamelCase			,       )
class    _UpperCAmelCase							(		_lowerCamelCase     ):
      def __init__(		self      ,       a__       ):
         super().__init__(a__       )
         A_     :				str			      =	config.num_labels
         A_     :				Union[str, Any]			      =	PoolFormerModel(a__       )

         # Final norm
         A_     :				str			      =	PoolFormerGroupNorm(config.hidden_sizes[-1]       )
         # Classifier head
         A_     :				Optional[int]			      =	(
             nn.Linear(config.hidden_sizes[-1]      ,       config.num_labels       ) if config.num_labels > 0 else nn.Identity()
         )

         # Initialize weights and apply final processing
         self.post_init()






      @add_start_docstrings_to_model_forward(a__       )
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT      ,       output_type=a__      ,       config_class=_CONFIG_FOR_DOC      ,       expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,       )
      def 		_lowerCamelCase      (		self      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       ):
         A_     :				List[Any]			      =	return_dict if return_dict is not None else self.config.use_return_dict

         A_     :				Optional[int]			      =	self.poolformer(
             a__      ,       output_hidden_states=a__      ,       return_dict=a__      ,       )

         A_     :				List[str]			      =	outputs[0]

         A_     :				int			      =	self.classifier(self.norm(a__       ).mean([-2, -1]       )       )

         A_     :				Dict			      =	None
         if labels is not None:
            if self.config.problem_type is None:
               if self.num_labels == 1:
                  A_     :				Tuple			      =	"""regression"""
               elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                  A_     :				Tuple			      =	"""single_label_classification"""
               else:
                  A_     :				Tuple			      =	"""multi_label_classification"""

            if self.config.problem_type == "regression":
               A_     :				int			      =	MSELoss()
               if self.num_labels == 1:
                  A_     :				List[str]			      =	loss_fct(logits.squeeze()      ,       labels.squeeze()       )
               else:
                  A_     :				Tuple			      =	loss_fct(a__      ,       a__       )
            elif self.config.problem_type == "single_label_classification":
               A_     :				str			      =	CrossEntropyLoss()
               A_     :				int			      =	loss_fct(logits.view(-1      ,       self.num_labels       )      ,       labels.view(-1       )       )
            elif self.config.problem_type == "multi_label_classification":
               A_     :				List[str]			      =	BCEWithLogitsLoss()
               A_     :				Optional[Any]			      =	loss_fct(a__      ,       a__       )

         if not return_dict:
            A_     :				Optional[Any]			      =	(logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

         return ImageClassifierOutputWithNoAttention(loss=a__      ,       logits=a__      ,       hidden_states=outputs.hidden_states       )













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    _UpperCAmelCase							:
      def __init__(		self      ,       a__      ,       a__=13      ,       a__=30      ,       a__=2      ,       a__=3      ,       a__=True      ,       a__=True      ,       a__=32      ,       a__=5      ,       a__=4      ,       a__=37      ,       a__="gelu"      ,       a__=0.1      ,       a__=0.1      ,       a__=10      ,       a__=0.02      ,       a__=3      ,       a__=0.6      ,       a__=None      ,       ):
         A_     :				int			      =	parent
         A_     :				Optional[int]			      =	batch_size
         A_     :				Any			      =	image_size
         A_     :				Optional[int]			      =	patch_size
         A_     :				int			      =	num_channels
         A_     :				str			      =	is_training
         A_     :				str			      =	use_labels
         A_     :				str			      =	hidden_size
         A_     :				Union[str, Any]			      =	num_hidden_layers
         A_     :				Tuple			      =	num_attention_heads
         A_     :				Any			      =	intermediate_size
         A_     :				List[Any]			      =	hidden_act
         A_     :				List[Any]			      =	hidden_dropout_prob
         A_     :				Optional[int]			      =	attention_probs_dropout_prob
         A_     :				str			      =	type_sequence_label_size
         A_     :				int			      =	initializer_range
         A_     :				List[Any]			      =	mask_ratio
         A_     :				str			      =	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[Any]			      =	(image_size // patch_size) ** 2
         A_     :				List[Any]			      =	int(math.ceil((1 - mask_ratio) * (num_patches + 1)       )       )
      def 		_lowerCamelCase      (		self       ):
         A_     :				List[str]			      =	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

         A_     :				int			      =	None
         if self.use_labels:
            A_     :				List[Any]			      =	ids_tensor([self.batch_size]      ,       self.type_sequence_label_size       )

         A_     :				Any			      =	self.get_config()

         return config, pixel_values, labels
      def 		_lowerCamelCase      (		self       ):
         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=a__      ,       initializer_range=self.initializer_range      ,       mask_ratio=self.mask_ratio      ,       )
      def 		_lowerCamelCase      (		self      ,       a__      ,       a__      ,       a__       ):
         A_     :				Optional[int]			      =	ViTMAEModel(config=a__       )
         model.to(a__       )
         model.eval()
         A_     :				int			      =	model(a__       )
         self.parent.assertEqual(result.last_hidden_state.shape      ,       (self.batch_size, self.seq_length, self.hidden_size)       )
      def 		_lowerCamelCase      (		self      ,       a__      ,       a__      ,       a__       ):
         A_     :				int			      =	ViTMAEForPreTraining(a__       )
         model.to(a__       )
         model.eval()
         A_     :				Optional[Any]			      =	model(a__       )
         A_     :				Dict			      =	(self.image_size // self.patch_size) ** 2
         A_     :				Optional[int]			      =	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_     :				Optional[int]			      =	1
         A_     :				Any			      =	ViTMAEForPreTraining(a__       )
         model.to(a__       )
         model.eval()
         A_     :				Dict			      =	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
         A_     :				List[Any]			      =	model(a__       )
         A_     :				Optional[int]			      =	self.patch_size**2
         self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, num_patches, expected_num_channels)       )






      def 		_lowerCamelCase      (		self       ):
         A_     :				Optional[Any]			      =	self.prepare_config_and_inputs()
         A_      ,				A_      ,				A_     :				Any			      =	config_and_inputs
         A_     :				Dict			      =	{"""pixel_values""": pixel_values}
         return config, inputs_dict






@require_torch
class    _UpperCAmelCase							(		_lowerCamelCase			,       _lowerCamelCase			,       unittest.TestCase     ):
      a						       =							(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
      a						       =							{'''feature-extraction''': ViTMAEModel} if is_torch_available() else {}

      a						       =							False
      a						       =							False
      a						       =							False
      a						       =							False
      def 		_lowerCamelCase      (		self       ):
         A_     :				int			      =	ViTMAEModelTester(self       )
         A_     :				List[Any]			      =	ConfigTester(self      ,       config_class=a__      ,       has_text_modality=a__      ,       hidden_size=37       )
      def 		_lowerCamelCase      (		self       ):
         self.config_tester.run_common_tests()
      @unittest.skip(reason="""ViTMAE does not use inputs_embeds"""       )
      def 		_lowerCamelCase      (		self       ):
         pass
      def 		_lowerCamelCase      (		self       ):
         A_      ,				A_     :				Union[str, Any]			      =	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            A_     :				Dict			      =	model_class(a__       )
            self.assertIsInstance(model.get_input_embeddings()      ,       (nn.Module)       )
            A_     :				int			      =	model.get_output_embeddings()
            self.assertTrue(x is None or isinstance(a__      ,       nn.Linear       )       )
      def 		_lowerCamelCase      (		self       ):
         A_      ,				A_     :				int			      =	self.model_tester.prepare_config_and_inputs_for_common()

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

            A_     :				List[str]			      =	["""pixel_values"""]
            self.assertListEqual(arg_names[:1]      ,       a__       )
      def 		_lowerCamelCase      (		self       ):
         A_     :				Dict			      =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a__       )
      def 		_lowerCamelCase      (		self       ):
         A_     :				Any			      =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_pretraining(*a__       )
      def 		_lowerCamelCase      (		self      ,       a__      ,       a__      ,       a__       ):
         # make masks reproducible
         np.random.seed(2       )

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

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

         super().check_pt_tf_models(a__      ,       a__      ,       a__       )
      def 		_lowerCamelCase      (		self       ):
         A_      ,				A_     :				int			      =	self.model_tester.prepare_config_and_inputs_for_common()

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

            A_     :				int			      =	outputs[0].cpu().numpy()
            A_     :				Union[str, Any]			      =	0

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

               # Make sure we don't have nans
               A_     :				Optional[int]			      =	after_outputs[0].cpu().numpy()
               A_     :				str			      =	0
               A_     :				Optional[int]			      =	np.amax(np.abs(out_a - out_a       )       )
               self.assertLessEqual(a__      ,       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 		_lowerCamelCase      (		self       ):
         pass
      @unittest.skip(
          reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""       )
      def 		_lowerCamelCase      (		self       ):
         pass
      @unittest.skip(
          reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""       )
      def 		_lowerCamelCase      (		self       ):
         pass
      @unittest.skip(reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"""       )
      def 		_lowerCamelCase      (		self       ):
         pass
      @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""       )
      def 		_lowerCamelCase      (		self       ):
         pass






      @slow
      def 		_lowerCamelCase      (		self       ):
         for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            A_     :				Union[str, Any]			      =	ViTMAEModel.from_pretrained(a__       )
            self.assertIsNotNone(a__       )






def     _lowerCAmelCase		(  ):
   '''simple docstring'''


   A_     :				List[Any]			      =	Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""     )
   return image







@require_torch
@require_vision
class    _UpperCAmelCase							(		unittest.TestCase     ):
      @cached_property
      def 		_lowerCamelCase      (		self       ):
         return ViTImageProcessor.from_pretrained("""facebook/vit-mae-base"""       ) if is_vision_available() else None






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

         A_     :				Union[str, Any]			      =	ViTMAEForPreTraining.from_pretrained("""facebook/vit-mae-base"""       ).to(a__       )

         A_     :				Optional[Any]			      =	self.default_image_processor
         A_     :				Union[str, Any]			      =	prepare_img()
         A_     :				Union[str, Any]			      =	image_processor(images=a__      ,       return_tensors="""pt"""       ).to(a__       )

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

         # forward pass
         with torch.no_grad():
            A_     :				Dict			      =	model(**a__      ,       noise=torch.from_numpy(a__       ).to(device=a__       )       )

         # verify the logits
         A_     :				Tuple			      =	torch.Size((1, 196, 768)       )
         self.assertEqual(outputs.logits.shape      ,       a__       )

         A_     :				List[str]			      =	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(a__       )      ,       atol=1E-4       )       )






