
"""simple docstring"""


from typing import Optional

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

from ...activations import ACTaFN
from ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


A_    							=					logging.get_logger(__name__)

# General docstring
A_    							=					"""ResNetConfig"""

# Base docstring
A_    							=					"""microsoft/resnet-50"""
A_    							=					[1, 2048, 7, 7]

# Image classification docstring
A_    							=					"""microsoft/resnet-50"""
A_    							=					"""tiger cat"""

A_    							=					[
    """microsoft/resnet-50""",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]




class  __lowerCamelCase  ( nn.Module			):




				def __init__(				self							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase = 3							,						UpperCAmelCase = 1							,						UpperCAmelCase = "relu"   ):
							super().__init__()
							lowerCamelCase_					=  nn.Convad(
							    UpperCAmelCase							,						UpperCAmelCase							,						kernel_size=UpperCAmelCase							,						stride=UpperCAmelCase							,						padding=kernel_size // 2							,						bias=UpperCAmelCase   )
							lowerCamelCase_					=  nn.BatchNormad(UpperCAmelCase   )
							lowerCamelCase_					=  ACTaFN[activation] if activation is not None else nn.Identity()





				def        UpperCAmelCase__						(				self							,						UpperCAmelCase   ):
							lowerCamelCase_					=  self.convolution(UpperCAmelCase   )
							lowerCamelCase_					=  self.normalization(UpperCAmelCase   )
							lowerCamelCase_					=  self.activation(UpperCAmelCase   )
							return hidden_state




class  __lowerCamelCase  ( nn.Module			):




				def __init__(				self							,						UpperCAmelCase   ):
							super().__init__()
							lowerCamelCase_					=  ResNetConvLayer(
							    config.num_channels							,						config.embedding_size							,						kernel_size=7							,						stride=2							,						activation=config.hidden_act   )
							lowerCamelCase_					=  nn.MaxPoolad(kernel_size=3							,						stride=2							,						padding=1   )
							lowerCamelCase_					=  config.num_channels





				def        UpperCAmelCase__						(				self							,						UpperCAmelCase   ):
							lowerCamelCase_					=  pixel_values.shape[1]
							if num_channels != self.num_channels:
										raise ValueError(
										    '''Make sure that the channel dimension of the pixel values match with the one set in the configuration.'''   )
							lowerCamelCase_					=  self.embedder(UpperCAmelCase   )
							lowerCamelCase_					=  self.pooler(UpperCAmelCase   )
							return embedding




class  __lowerCamelCase  ( nn.Module			):




				def __init__(				self							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase = 2   ):
							super().__init__()
							lowerCamelCase_					=  nn.Convad(UpperCAmelCase							,						UpperCAmelCase							,						kernel_size=1							,						stride=UpperCAmelCase							,						bias=UpperCAmelCase   )
							lowerCamelCase_					=  nn.BatchNormad(UpperCAmelCase   )





				def        UpperCAmelCase__						(				self							,						UpperCAmelCase   ):
							lowerCamelCase_					=  self.convolution(UpperCAmelCase   )
							lowerCamelCase_					=  self.normalization(UpperCAmelCase   )
							return hidden_state




class  __lowerCamelCase  ( nn.Module			):




				def __init__(				self							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase = 1							,						UpperCAmelCase = "relu"   ):
							super().__init__()
							lowerCamelCase_					=  in_channels != out_channels or stride != 1
							lowerCamelCase_					=  (
							    ResNetShortCut(UpperCAmelCase							,						UpperCAmelCase							,						stride=UpperCAmelCase   ) if should_apply_shortcut else nn.Identity()
							)
							lowerCamelCase_					=  nn.Sequential(
							    ResNetConvLayer(UpperCAmelCase							,						UpperCAmelCase							,						stride=UpperCAmelCase   )							,						ResNetConvLayer(UpperCAmelCase							,						UpperCAmelCase							,						activation=UpperCAmelCase   )							,						)
							lowerCamelCase_					=  ACTaFN[activation]





				def        UpperCAmelCase__						(				self							,						UpperCAmelCase   ):
							lowerCamelCase_					=  hidden_state
							lowerCamelCase_					=  self.layer(UpperCAmelCase   )
							lowerCamelCase_					=  self.shortcut(UpperCAmelCase   )
							hidden_state += residual
							lowerCamelCase_					=  self.activation(UpperCAmelCase   )
							return hidden_state




class  __lowerCamelCase  ( nn.Module			):




				def __init__(				self							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase = 1							,						UpperCAmelCase = "relu"							,						UpperCAmelCase = 4   ):
							super().__init__()
							lowerCamelCase_					=  in_channels != out_channels or stride != 1
							lowerCamelCase_					=  out_channels // reduction
							lowerCamelCase_					=  (
							    ResNetShortCut(UpperCAmelCase							,						UpperCAmelCase							,						stride=UpperCAmelCase   ) if should_apply_shortcut else nn.Identity()
							)
							lowerCamelCase_					=  nn.Sequential(
							    ResNetConvLayer(UpperCAmelCase							,						UpperCAmelCase							,						kernel_size=1   )							,						ResNetConvLayer(UpperCAmelCase							,						UpperCAmelCase							,						stride=UpperCAmelCase   )							,						ResNetConvLayer(UpperCAmelCase							,						UpperCAmelCase							,						kernel_size=1							,						activation=UpperCAmelCase   )							,						)
							lowerCamelCase_					=  ACTaFN[activation]





				def        UpperCAmelCase__						(				self							,						UpperCAmelCase   ):
							lowerCamelCase_					=  hidden_state
							lowerCamelCase_					=  self.layer(UpperCAmelCase   )
							lowerCamelCase_					=  self.shortcut(UpperCAmelCase   )
							hidden_state += residual
							lowerCamelCase_					=  self.activation(UpperCAmelCase   )
							return hidden_state




class  __lowerCamelCase  ( nn.Module			):




				def __init__(				self							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase = 2							,						UpperCAmelCase = 2							,						):
							super().__init__()

							lowerCamelCase_					=  ResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else ResNetBasicLayer

							lowerCamelCase_					=  nn.Sequential(
							    # downsampling is done in the first layer with stride of 2
							    layer(UpperCAmelCase							,						UpperCAmelCase							,						stride=UpperCAmelCase							,						activation=config.hidden_act   )							,						*[layer(UpperCAmelCase							,						UpperCAmelCase							,						activation=config.hidden_act   ) for _ in range(depth - 1   )]							,						)





				def        UpperCAmelCase__						(				self							,						UpperCAmelCase   ):
							lowerCamelCase_					=  input
							for layer in self.layers:
										lowerCamelCase_					=  layer(UpperCAmelCase   )
							return hidden_state




class  __lowerCamelCase  ( nn.Module			):




				def __init__(				self							,						UpperCAmelCase   ):
							super().__init__()
							lowerCamelCase_					=  nn.ModuleList([]   )
							# based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
							self.stages.append(
							    ResNetStage(
							        UpperCAmelCase							,						config.embedding_size							,						config.hidden_sizes[0]							,						stride=2 if config.downsample_in_first_stage else 1							,						depth=config.depths[0]							,						)   )
							lowerCamelCase_					=  zip(config.hidden_sizes							,						config.hidden_sizes[1:]   )
							for (in_channels, out_channels), depth in zip(UpperCAmelCase							,						config.depths[1:]   ):
										self.stages.append(ResNetStage(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						depth=UpperCAmelCase   )   )





				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase = False							,						UpperCAmelCase = True   ):
							lowerCamelCase_					=  () if output_hidden_states else None

							for stage_module in self.stages:
										if output_hidden_states:
													lowerCamelCase_					=  hidden_states + (hidden_state,)

										lowerCamelCase_					=  stage_module(UpperCAmelCase   )

							if output_hidden_states:
										lowerCamelCase_					=  hidden_states + (hidden_state,)

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

							return BaseModelOutputWithNoAttention(
							    last_hidden_state=UpperCAmelCase							,						hidden_states=UpperCAmelCase							,						)




class  __lowerCamelCase  ( lowerCAmelCase			):
				a__:      Tuple          =    ResNetConfig
				a__:      Tuple          =    'resnet'
				a__:      List[Any]          =    'pixel_values'
				a__:      Union[str, Any]          =    True




				def        UpperCAmelCase__						(				self							,						UpperCAmelCase   ):
							if isinstance(UpperCAmelCase							,						nn.Convad   ):
										nn.init.kaiming_normal_(module.weight							,						mode='''fan_out'''							,						nonlinearity='''relu'''   )
							elif isinstance(UpperCAmelCase							,						(nn.BatchNormad, nn.GroupNorm)   ):
										nn.init.constant_(module.weight							,						1   )
										nn.init.constant_(module.bias							,						0   )





				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase=False   ):
							if isinstance(UpperCAmelCase							,						UpperCAmelCase   ):
										lowerCamelCase_					=  value

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

    Parameters:
        config ([`ResNetConfig`]): 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.
"""

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

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""




@add_start_docstrings(
    'The bare ResNet model outputting raw features without any specific head on top.'						,	lowerCAmelCase						,	)
class  __lowerCamelCase  ( lowerCAmelCase			):




				def __init__(				self							,						UpperCAmelCase   ):
							super().__init__(UpperCAmelCase   )
							lowerCamelCase_					=  config
							lowerCamelCase_					=  ResNetEmbeddings(UpperCAmelCase   )
							lowerCamelCase_					=  ResNetEncoder(UpperCAmelCase   )
							lowerCamelCase_					=  nn.AdaptiveAvgPoolad((1, 1)   )
							# Initialize weights and apply final processing
							self.post_init()





				@add_start_docstrings_to_model_forward(UpperCAmelCase   )
				@add_code_sample_docstrings(
				    checkpoint=_CHECKPOINT_FOR_DOC							,						output_type=UpperCAmelCase							,						config_class=_CONFIG_FOR_DOC							,						modality='''vision'''							,						expected_output=_EXPECTED_OUTPUT_SHAPE							,						)
				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase = None							,						UpperCAmelCase = None   ):
							lowerCamelCase_					=  (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							lowerCamelCase_					=  return_dict if return_dict is not None else self.config.use_return_dict

							lowerCamelCase_					=  self.embedder(UpperCAmelCase   )

							lowerCamelCase_					=  self.encoder(
							    UpperCAmelCase							,						output_hidden_states=UpperCAmelCase							,						return_dict=UpperCAmelCase   )

							lowerCamelCase_					=  encoder_outputs[0]

							lowerCamelCase_					=  self.pooler(UpperCAmelCase   )

							if not return_dict:
										return (last_hidden_state, pooled_output) + encoder_outputs[1:]

							return BaseModelOutputWithPoolingAndNoAttention(
							    last_hidden_state=UpperCAmelCase							,						pooler_output=UpperCAmelCase							,						hidden_states=encoder_outputs.hidden_states							,						)




@add_start_docstrings(
    '\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    '						,	lowerCAmelCase						,	)
class  __lowerCamelCase  ( lowerCAmelCase			):




				def __init__(				self							,						UpperCAmelCase   ):
							super().__init__(UpperCAmelCase   )
							lowerCamelCase_					=  config.num_labels
							lowerCamelCase_					=  ResNetModel(UpperCAmelCase   )
							# classification head
							lowerCamelCase_					=  nn.Sequential(
							    nn.Flatten()							,						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(UpperCAmelCase   )
				@add_code_sample_docstrings(
				    checkpoint=_IMAGE_CLASS_CHECKPOINT							,						output_type=UpperCAmelCase							,						config_class=_CONFIG_FOR_DOC							,						expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							,						)
				def        UpperCAmelCase__						(				self							,						UpperCAmelCase = None							,						UpperCAmelCase = None							,						UpperCAmelCase = None							,						UpperCAmelCase = None							,						):
							lowerCamelCase_					=  return_dict if return_dict is not None else self.config.use_return_dict

							lowerCamelCase_					=  self.resnet(UpperCAmelCase							,						output_hidden_states=UpperCAmelCase							,						return_dict=UpperCAmelCase   )

							lowerCamelCase_					=  outputs.pooler_output if return_dict else outputs[1]

							lowerCamelCase_					=  self.classifier(UpperCAmelCase   )

							lowerCamelCase_					=  None

							if labels is not None:
										if self.config.problem_type is None:
													if self.num_labels == 1:
																lowerCamelCase_					=  '''regression'''
													elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																lowerCamelCase_					=  '''single_label_classification'''
													else:
																lowerCamelCase_					=  '''multi_label_classification'''
										if self.config.problem_type == "regression":
													lowerCamelCase_					=  MSELoss()
													if self.num_labels == 1:
																lowerCamelCase_					=  loss_fct(logits.squeeze()							,						labels.squeeze()   )
													else:
																lowerCamelCase_					=  loss_fct(UpperCAmelCase							,						UpperCAmelCase   )
										elif self.config.problem_type == "single_label_classification":
													lowerCamelCase_					=  CrossEntropyLoss()
													lowerCamelCase_					=  loss_fct(logits.view(-1							,						self.num_labels   )							,						labels.view(-1   )   )
										elif self.config.problem_type == "multi_label_classification":
													lowerCamelCase_					=  BCEWithLogitsLoss()
													lowerCamelCase_					=  loss_fct(UpperCAmelCase							,						UpperCAmelCase   )

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

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




@add_start_docstrings(
    '\n    ResNet backbone, to be used with frameworks like DETR and MaskFormer.\n    '						,	lowerCAmelCase						,	)
class  __lowerCamelCase  ( lowerCAmelCase						,	lowerCAmelCase			):




				def __init__(				self							,						UpperCAmelCase   ):
							super().__init__(UpperCAmelCase   )
							super()._init_backbone(UpperCAmelCase   )

							lowerCamelCase_					=  [config.embedding_size] + config.hidden_sizes
							lowerCamelCase_					=  ResNetEmbeddings(UpperCAmelCase   )
							lowerCamelCase_					=  ResNetEncoder(UpperCAmelCase   )

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





				@add_start_docstrings_to_model_forward(UpperCAmelCase   )
				@replace_return_docstrings(output_type=UpperCAmelCase							,						config_class=_CONFIG_FOR_DOC   )
				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase = None							,						UpperCAmelCase = None   ):
							lowerCamelCase_					=  return_dict if return_dict is not None else self.config.use_return_dict
							lowerCamelCase_					=  (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)

							lowerCamelCase_					=  self.embedder(UpperCAmelCase   )

							lowerCamelCase_					=  self.encoder(UpperCAmelCase							,						output_hidden_states=UpperCAmelCase							,						return_dict=UpperCAmelCase   )

							lowerCamelCase_					=  outputs.hidden_states

							lowerCamelCase_					=  ()
							for idx, stage in enumerate(self.stage_names   ):
										if stage in self.out_features:
													feature_maps += (hidden_states[idx],)

							if not return_dict:
										lowerCamelCase_					=  (feature_maps,)
										if output_hidden_states:
													output += (outputs.hidden_states,)
										return output

							return BackboneOutput(
							    feature_maps=UpperCAmelCase							,						hidden_states=outputs.hidden_states if output_hidden_states else None							,						attentions=UpperCAmelCase							,						)







'''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		)		)

