




'''simple docstring'''





from typing import Optional, Union

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

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithPoolingAndNoAttention, 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_mobilenet_va import MobileNetVaConfig


_lowerCamelCase      :  Optional[Any]      						=						logging.get_logger(__name__)


# General docstring
_lowerCamelCase      :  Tuple      						=						"MobileNetV1Config"

# Base docstring
_lowerCamelCase      :  Any      						=						"google/mobilenet_v1_1.0_224"
_lowerCamelCase      :  int      						=						[1, 1_0_2_4, 7, 7]

# Image classification docstring
_lowerCamelCase      :  Dict      						=						"google/mobilenet_v1_1.0_224"
_lowerCamelCase      :  int      						=						"tabby, tabby cat"


_lowerCamelCase      :  List[str]      						=						[
    "google/mobilenet_v1_1.0_224",
    "google/mobilenet_v1_0.75_192",
    # See all MobileNetV1 models at https://huggingface.co/models?filter=mobilenet_v1
]




def 							_UpperCAmelCase  (UpperCamelCase_ : int      ,    UpperCamelCase_ : int      ,    UpperCamelCase_ : int=None       ):
				'''simple docstring'''
				_lowerCAmelCase    :    Any										=      {}

				if isinstance(UpperCamelCase_      ,    UpperCamelCase_       ):
								_lowerCAmelCase    :    Optional[int]										=      model.mobilenet_va
				else:
								_lowerCAmelCase    :    Any										=      model

				_lowerCAmelCase    :    List[Any]										=      """MobilenetV1/Conv2d_0/"""
				_lowerCAmelCase    :    Dict										=      backbone.conv_stem.convolution.weight
				_lowerCAmelCase    :    str										=      backbone.conv_stem.normalization.bias
				_lowerCAmelCase    :    List[Any]										=      backbone.conv_stem.normalization.weight
				_lowerCAmelCase    :    List[str]										=      backbone.conv_stem.normalization.running_mean
				_lowerCAmelCase    :    Union[str, Any]										=      backbone.conv_stem.normalization.running_var

				for i in range(13       ):
								_lowerCAmelCase    :    Tuple										=      i + 1
								_lowerCAmelCase    :    List[Any]										=      i * 2

								_lowerCAmelCase    :    int										=      backbone.layer[pt_index]
								_lowerCAmelCase    :    Dict										=      F"MobilenetV1/Conv2d_{tf_index}_depthwise/"
								_lowerCAmelCase    :    int										=      pointer.convolution.weight
								_lowerCAmelCase    :    List[str]										=      pointer.normalization.bias
								_lowerCAmelCase    :    Union[str, Any]										=      pointer.normalization.weight
								_lowerCAmelCase    :    Optional[int]										=      pointer.normalization.running_mean
								_lowerCAmelCase    :    Optional[int]										=      pointer.normalization.running_var

								_lowerCAmelCase    :    Optional[int]										=      backbone.layer[pt_index + 1]
								_lowerCAmelCase    :    Union[str, Any]										=      F"MobilenetV1/Conv2d_{tf_index}_pointwise/"
								_lowerCAmelCase    :    Union[str, Any]										=      pointer.convolution.weight
								_lowerCAmelCase    :    List[Any]										=      pointer.normalization.bias
								_lowerCAmelCase    :    Optional[int]										=      pointer.normalization.weight
								_lowerCAmelCase    :    int										=      pointer.normalization.running_mean
								_lowerCAmelCase    :    Optional[Any]										=      pointer.normalization.running_var

				if isinstance(UpperCamelCase_      ,    UpperCamelCase_       ):
								_lowerCAmelCase    :    Union[str, Any]										=      """MobilenetV1/Logits/Conv2d_1c_1x1/"""
								_lowerCAmelCase    :    Optional[Any]										=      model.classifier.weight
								_lowerCAmelCase    :    Optional[int]										=      model.classifier.bias

				return tf_to_pt_map




def 							_UpperCAmelCase  (UpperCamelCase_ : Tuple      ,    UpperCamelCase_ : int      ,    UpperCamelCase_ : Optional[Any]       ):
				'''simple docstring'''
				try:
								import numpy as np
								import tensorflow as tf
				except ImportError:
								logger.error(
								    """Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see """
								    """https://www.tensorflow.org/install/ for installation instructions."""       )
								raise

				# Load weights from TF model
				_lowerCAmelCase    :    Optional[Any]										=      tf.train.list_variables(UpperCamelCase_       )
				_lowerCAmelCase    :    int										=      {}
				for name, shape in init_vars:
								logger.info(F"Loading TF weight {name} with shape {shape}"       )
								_lowerCAmelCase    :    Optional[Any]										=      tf.train.load_variable(UpperCamelCase_      ,    UpperCamelCase_       )
								_lowerCAmelCase    :    Optional[Any]										=      array

				# Build TF to PyTorch weights loading map
				_lowerCAmelCase    :    Union[str, Any]										=      _build_tf_to_pytorch_map(UpperCamelCase_      ,    UpperCamelCase_      ,    UpperCamelCase_       )

				for name, pointer in tf_to_pt_map.items():
								logger.info(F"Importing {name}"       )
								if name not in tf_weights:
												logger.info(F"{name} not in tf pre-trained weights, skipping"       )
												continue

								_lowerCAmelCase    :    str										=      tf_weights[name]

								if "depthwise_weights" in name:
												logger.info("""Transposing depthwise"""       )
												_lowerCAmelCase    :    Union[str, Any]										=      np.transpose(UpperCamelCase_      ,    (2, 3, 0, 1)       )
								elif "weights" in name:
												logger.info("""Transposing"""       )
												if len(pointer.shape       ) == 2:  # copying into linear layer
																_lowerCAmelCase    :    Any										=      array.squeeze().transpose()
												else:
																_lowerCAmelCase    :    List[str]										=      np.transpose(UpperCamelCase_      ,    (3, 2, 0, 1)       )

								if pointer.shape != array.shape:
												raise ValueError(F"Pointer shape {pointer.shape} and array shape {array.shape} mismatched"       )

								logger.info(F"Initialize PyTorch weight {name} {array.shape}"       )
								_lowerCAmelCase    :    Dict										=      torch.from_numpy(UpperCamelCase_       )

								tf_weights.pop(UpperCamelCase_      ,    UpperCamelCase_       )
								tf_weights.pop(name + """/RMSProp"""      ,    UpperCamelCase_       )
								tf_weights.pop(name + """/RMSProp_1"""      ,    UpperCamelCase_       )
								tf_weights.pop(name + """/ExponentialMovingAverage"""      ,    UpperCamelCase_       )

				logger.info(F"Weights not copied to PyTorch model: {', '.join(tf_weights.keys()       )}"       )
				return model




def 							_UpperCAmelCase  (UpperCamelCase_ : torch.Tensor      ,    UpperCamelCase_ : nn.Convad       ):
				'''simple docstring'''
				_lowerCAmelCase    :    List[str]										=      features.shape[-2:]
				_lowerCAmelCase    :    str										=      conv_layer.stride
				_lowerCAmelCase    :    int										=      conv_layer.kernel_size

				if in_height % stride_height == 0:
								_lowerCAmelCase    :    int										=      max(kernel_height - stride_height      ,    0       )
				else:
								_lowerCAmelCase    :    List[str]										=      max(kernel_height - (in_height % stride_height)      ,    0       )

				if in_width % stride_width == 0:
								_lowerCAmelCase    :    Dict										=      max(kernel_width - stride_width      ,    0       )
				else:
								_lowerCAmelCase    :    Union[str, Any]										=      max(kernel_width - (in_width % stride_width)      ,    0       )

				_lowerCAmelCase    :    List[Any]										=      pad_along_width // 2
				_lowerCAmelCase    :    Dict										=      pad_along_width - pad_left
				_lowerCAmelCase    :    Optional[int]										=      pad_along_height // 2
				_lowerCAmelCase    :    Optional[int]										=      pad_along_height - pad_top

				_lowerCAmelCase    :    Any										=      (pad_left, pad_right, pad_top, pad_bottom)
				return nn.functional.pad(UpperCamelCase_      ,    UpperCamelCase_      ,    """constant"""      ,    0.0       )
class        __snake_case  (nn.Module		):



				def __init__(      self      :		List[Any]		,	_UpperCAmelCase      :		MobileNetVaConfig		,	_UpperCAmelCase      :		int		,	_UpperCAmelCase      :		int		,	_UpperCAmelCase      :		int		,	_UpperCAmelCase      :		Optional[int] = 1		,	_UpperCAmelCase      :		Optional[int] = 1		,	_UpperCAmelCase      :		bool = False		,	_UpperCAmelCase      :		Optional[bool] = True		,	_UpperCAmelCase      :		Optional[bool or str] = True		,	)  ->    None:





								'''simple docstring'''




								super().__init__()
								_lowerCAmelCase    :    Optional[Any]										=      config

								if in_channels % groups != 0:
												raise ValueError(f"Input channels ({in_channels}) are not divisible by {groups} groups."      )
								if out_channels % groups != 0:
												raise ValueError(f"Output channels ({out_channels}) are not divisible by {groups} groups."      )

								_lowerCAmelCase    :    Union[str, Any]										=      0 if config.tf_padding else int((kernel_size - 1) / 2      )

								_lowerCAmelCase    :    Dict										=      nn.Convad(
								    in_channels=_UpperCAmelCase		,	out_channels=_UpperCAmelCase		,	kernel_size=_UpperCAmelCase		,	stride=_UpperCAmelCase		,	padding=_UpperCAmelCase		,	groups=_UpperCAmelCase		,	bias=_UpperCAmelCase		,	padding_mode="""zeros"""		,	)

								if use_normalization:
												_lowerCAmelCase    :    List[Any]										=      nn.BatchNormad(
												    num_features=_UpperCAmelCase		,	eps=config.layer_norm_eps		,	momentum=0.9997		,	affine=_UpperCAmelCase		,	track_running_stats=_UpperCAmelCase		,	)
								else:
												_lowerCAmelCase    :    Optional[int]										=      None

								if use_activation:
												if isinstance(_UpperCAmelCase		,	_UpperCAmelCase      ):
																_lowerCAmelCase    :    int										=      ACTaFN[use_activation]
												elif isinstance(config.hidden_act		,	_UpperCAmelCase      ):
																_lowerCAmelCase    :    str										=      ACTaFN[config.hidden_act]
												else:
																_lowerCAmelCase    :    Dict										=      config.hidden_act
								else:
												_lowerCAmelCase    :    List[Any]										=      None



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]		,	_UpperCAmelCase      :		torch.Tensor      )  ->    torch.Tensor:





								'''simple docstring'''




								if self.config.tf_padding:
												_lowerCAmelCase    :    int										=      apply_tf_padding(_UpperCAmelCase		,	self.convolution      )
								_lowerCAmelCase    :    Union[str, Any]										=      self.convolution(_UpperCAmelCase      )
								if self.normalization is not None:
												_lowerCAmelCase    :    Tuple										=      self.normalization(_UpperCAmelCase      )
								if self.activation is not None:
												_lowerCAmelCase    :    Union[str, Any]										=      self.activation(_UpperCAmelCase      )
								return features






class        __snake_case  (_a		):
				lowerCAmelCase__				=     MobileNetVaConfig
				lowerCAmelCase__				=     load_tf_weights_in_mobilenet_va
				lowerCAmelCase__				=     "mobilenet_v1"
				lowerCAmelCase__				=     "pixel_values"
				lowerCAmelCase__				=     False



				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[int]		,	_UpperCAmelCase      :		Union[nn.Linear, nn.Convad]      )  ->    None:





								'''simple docstring'''




								if isinstance(_UpperCAmelCase		,	(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(_UpperCAmelCase		,	nn.BatchNormad      ):
												module.bias.data.zero_()
												module.weight.data.fill_(1.0      )






_lowerCamelCase      :  Union[str, Any]      						=						R"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`MobileNetV1Config`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

_lowerCamelCase      :  List[Any]      						=						R"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`MobileNetV1ImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"
@add_start_docstrings(
    "The bare MobileNetV1 model outputting raw hidden-states without any specific head on top."						,	_a						,	)
class        __snake_case  (_a		):



				def __init__(      self      :		Optional[int]		,	_UpperCAmelCase      :		MobileNetVaConfig		,	_UpperCAmelCase      :		bool = True      )  ->    Optional[Any]:





								'''simple docstring'''




								super().__init__(_UpperCAmelCase      )
								_lowerCAmelCase    :    Dict										=      config

								_lowerCAmelCase    :    Optional[Any]										=      32
								_lowerCAmelCase    :    Dict										=      max(int(depth * config.depth_multiplier      )		,	config.min_depth      )

								_lowerCAmelCase    :    int										=      MobileNetVaConvLayer(
								    _UpperCAmelCase		,	in_channels=config.num_channels		,	out_channels=_UpperCAmelCase		,	kernel_size=3		,	stride=2		,	)

								_lowerCAmelCase    :    Tuple										=      [1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

								_lowerCAmelCase    :    int										=      nn.ModuleList()
								for i in range(13      ):
												_lowerCAmelCase    :    List[Any]										=      out_channels

												if strides[i] == 2 or i == 0:
																depth *= 2
																_lowerCAmelCase    :    str										=      max(int(depth * config.depth_multiplier      )		,	config.min_depth      )

												self.layer.append(
												    MobileNetVaConvLayer(
												        _UpperCAmelCase		,	in_channels=_UpperCAmelCase		,	out_channels=_UpperCAmelCase		,	kernel_size=3		,	stride=strides[i]		,	groups=_UpperCAmelCase		,	)      )

												self.layer.append(
												    MobileNetVaConvLayer(
												        _UpperCAmelCase		,	in_channels=_UpperCAmelCase		,	out_channels=_UpperCAmelCase		,	kernel_size=1		,	)      )

								_lowerCAmelCase    :    List[Any]										=      nn.AdaptiveAvgPoolad((1, 1)      ) if add_pooling_layer else None

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



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





								'''simple docstring'''




								raise NotImplementedError



				@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       SCREAMING_SNAKE_CASE       (      self      :		Any		,	_UpperCAmelCase      :		Optional[torch.Tensor] = None		,	_UpperCAmelCase      :		Optional[bool] = None		,	_UpperCAmelCase      :		Optional[bool] = None		,	)  ->    Union[tuple, BaseModelOutputWithPoolingAndNoAttention]:





								'''simple docstring'''




								_lowerCAmelCase    :    Tuple										=      (
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								_lowerCAmelCase    :    Optional[Any]										=      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"""      )

								_lowerCAmelCase    :    List[Any]										=      self.conv_stem(_UpperCAmelCase      )

								_lowerCAmelCase    :    int										=      () if output_hidden_states else None

								for i, layer_module in enumerate(self.layer      ):
												_lowerCAmelCase    :    List[str]										=      layer_module(_UpperCAmelCase      )

												if output_hidden_states:
																_lowerCAmelCase    :    List[Any]										=      all_hidden_states + (hidden_states,)

								_lowerCAmelCase    :    Union[str, Any]										=      hidden_states

								if self.pooler is not None:
												_lowerCAmelCase    :    Tuple										=      torch.flatten(self.pooler(_UpperCAmelCase      )		,	start_dim=1      )
								else:
												_lowerCAmelCase    :    List[str]										=      None

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

								return BaseModelOutputWithPoolingAndNoAttention(
								    last_hidden_state=_UpperCAmelCase		,	pooler_output=_UpperCAmelCase		,	hidden_states=_UpperCAmelCase		,	)






@add_start_docstrings(
    "\n    MobileNetV1 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "						,	_a						,	)
class        __snake_case  (_a		):



				def __init__(      self      :		Union[str, Any]		,	_UpperCAmelCase      :		MobileNetVaConfig      )  ->    None:





								'''simple docstring'''




								super().__init__(_UpperCAmelCase      )

								_lowerCAmelCase    :    Any										=      config.num_labels
								_lowerCAmelCase    :    int										=      MobileNetVaModel(_UpperCAmelCase      )

								_lowerCAmelCase    :    List[Any]										=      self.mobilenet_va.layer[-1].convolution.out_channels

								# Classifier head
								_lowerCAmelCase    :    Optional[Any]										=      nn.Dropout(config.classifier_dropout_prob		,	inplace=_UpperCAmelCase      )
								_lowerCAmelCase    :    str										=      nn.Linear(_UpperCAmelCase		,	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       SCREAMING_SNAKE_CASE       (      self      :		List[Any]		,	_UpperCAmelCase      :		Optional[torch.Tensor] = None		,	_UpperCAmelCase      :		Optional[bool] = None		,	_UpperCAmelCase      :		Optional[torch.Tensor] = None		,	_UpperCAmelCase      :		Optional[bool] = None		,	)  ->    Union[tuple, ImageClassifierOutputWithNoAttention]:





								'''simple docstring'''




								_lowerCAmelCase    :    List[Any]										=      return_dict if return_dict is not None else self.config.use_return_dict

								_lowerCAmelCase    :    Union[str, Any]										=      self.mobilenet_va(_UpperCAmelCase		,	output_hidden_states=_UpperCAmelCase		,	return_dict=_UpperCAmelCase      )

								_lowerCAmelCase    :    Optional[int]										=      outputs.pooler_output if return_dict else outputs[1]

								_lowerCAmelCase    :    List[str]										=      self.classifier(self.dropout(_UpperCAmelCase      )      )

								_lowerCAmelCase    :    Dict										=      None
								if labels is not None:
												if self.config.problem_type is None:
																if self.num_labels == 1:
																				_lowerCAmelCase    :    Optional[int]										=      """regression"""
																elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																				_lowerCAmelCase    :    str										=      """single_label_classification"""
																else:
																				_lowerCAmelCase    :    Union[str, Any]										=      """multi_label_classification"""

												if self.config.problem_type == "regression":
																_lowerCAmelCase    :    Any										=      MSELoss()
																if self.num_labels == 1:
																				_lowerCAmelCase    :    Union[str, Any]										=      loss_fct(logits.squeeze()		,	labels.squeeze()      )
																else:
																				_lowerCAmelCase    :    Union[str, Any]										=      loss_fct(_UpperCAmelCase		,	_UpperCAmelCase      )
												elif self.config.problem_type == "single_label_classification":
																_lowerCAmelCase    :    Optional[Any]										=      CrossEntropyLoss()
																_lowerCAmelCase    :    Dict										=      loss_fct(logits.view(-1		,	self.num_labels      )		,	labels.view(-1      )      )
												elif self.config.problem_type == "multi_label_classification":
																_lowerCAmelCase    :    Optional[Any]										=      BCEWithLogitsLoss()
																_lowerCAmelCase    :    Any										=      loss_fct(_UpperCAmelCase		,	_UpperCAmelCase      )

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

								return ImageClassifierOutputWithNoAttention(
								    loss=_UpperCAmelCase		,	logits=_UpperCAmelCase		,	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        __snake_case  :



				def __init__(      self      :		Tuple		,	_UpperCAmelCase      :		Dict		,	_UpperCAmelCase      :		int=13		,	_UpperCAmelCase      :		Union[str, Any]=30		,	_UpperCAmelCase      :		List[str]=2		,	_UpperCAmelCase      :		Tuple=3		,	_UpperCAmelCase      :		Any=True		,	_UpperCAmelCase      :		Tuple=True		,	_UpperCAmelCase      :		int=32		,	_UpperCAmelCase      :		List[Any]=5		,	_UpperCAmelCase      :		int=4		,	_UpperCAmelCase      :		Union[str, Any]=37		,	_UpperCAmelCase      :		Optional[Any]="gelu"		,	_UpperCAmelCase      :		Dict=0.1		,	_UpperCAmelCase      :		Optional[int]=0.1		,	_UpperCAmelCase      :		List[Any]=10		,	_UpperCAmelCase      :		List[str]=0.02		,	_UpperCAmelCase      :		Optional[int]=3		,	_UpperCAmelCase      :		Any=0.6		,	_UpperCAmelCase      :		Any=None		,	)  ->    Tuple:





								'''simple docstring'''




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

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



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    List[str]										=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

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

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

								return config, pixel_values, labels



				def       SCREAMING_SNAKE_CASE       (      self      :		str      )  ->    Dict:





								'''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=_UpperCAmelCase		,	initializer_range=self.initializer_range		,	mask_ratio=self.mask_ratio		,	)



				def       SCREAMING_SNAKE_CASE       (      self      :		int		,	_UpperCAmelCase      :		str		,	_UpperCAmelCase      :		Optional[int]		,	_UpperCAmelCase      :		Dict      )  ->    Optional[int]:





								'''simple docstring'''




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



				def       SCREAMING_SNAKE_CASE       (      self      :		Any		,	_UpperCAmelCase      :		Optional[int]		,	_UpperCAmelCase      :		Optional[Any]		,	_UpperCAmelCase      :		Any      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[int]										=      ViTMAEForPreTraining(_UpperCAmelCase      )
								model.to(_UpperCAmelCase      )
								model.eval()
								_lowerCAmelCase    :    int										=      model(_UpperCAmelCase      )
								_lowerCAmelCase    :    Union[str, Any]										=      (self.image_size // self.patch_size) ** 2
								_lowerCAmelCase    :    Union[str, 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
								_lowerCAmelCase    :    Dict										=      1
								_lowerCAmelCase    :    Dict										=      ViTMAEForPreTraining(_UpperCAmelCase      )
								model.to(_UpperCAmelCase      )
								model.eval()
								_lowerCAmelCase    :    Union[str, Any]										=      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
								_lowerCAmelCase    :    str										=      model(_UpperCAmelCase      )
								_lowerCAmelCase    :    Optional[int]										=      self.patch_size**2
								self.parent.assertEqual(result.logits.shape		,	(self.batch_size, num_patches, expected_num_channels)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Any:





								'''simple docstring'''




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






@require_torch
class        __snake_case  (_a						,	_a						,	unittest.TestCase		):
				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       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    List[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    List[str]										=      ViTMAEModelTester(self      )
								_lowerCAmelCase    :    int										=      ConfigTester(self		,	config_class=_UpperCAmelCase		,	has_text_modality=_UpperCAmelCase		,	hidden_size=37      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple      )  ->    Optional[int]:





								'''simple docstring'''




								self.config_tester.run_common_tests()



				@unittest.skip(reason="""ViTMAE does not use inputs_embeds"""      )
				def       SCREAMING_SNAKE_CASE       (      self      :		str      )  ->    int:





								'''simple docstring'''




								pass



				def       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    str:





								'''simple docstring'''




								_lowerCAmelCase					,    _lowerCAmelCase    :    int										=      self.model_tester.prepare_config_and_inputs_for_common()

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



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Any:





								'''simple docstring'''




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

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

												_lowerCAmelCase    :    Any										=      ["""pixel_values"""]
												self.assertListEqual(arg_names[:1]		,	_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[Any]      )  ->    Optional[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[int]      )  ->    int:





								'''simple docstring'''




								_lowerCAmelCase    :    Any										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Dict		,	_UpperCAmelCase      :		Optional[int]		,	_UpperCAmelCase      :		Optional[Any]		,	_UpperCAmelCase      :		Optional[int]      )  ->    List[Any]:





								'''simple docstring'''




								np.random.seed(2      )

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

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

								super().check_pt_tf_models(_UpperCAmelCase		,	_UpperCAmelCase		,	_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		str      )  ->    Optional[int]:





								'''simple docstring'''




								_lowerCAmelCase					,    _lowerCAmelCase    :    str										=      self.model_tester.prepare_config_and_inputs_for_common()

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

												_lowerCAmelCase    :    List[Any]										=      outputs[0].cpu().numpy()
												_lowerCAmelCase    :    Union[str, Any]										=      0

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

																# Make sure we don't have nans
																_lowerCAmelCase    :    int										=      after_outputs[0].cpu().numpy()
																_lowerCAmelCase    :    Optional[Any]										=      0
																_lowerCAmelCase    :    List[Any]										=      np.amax(np.abs(out_a - out_a      )      )
																self.assertLessEqual(_UpperCAmelCase		,	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      :		int      )  ->    Tuple:





								'''simple docstring'''




								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      )  ->    str:





								'''simple docstring'''




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





								'''simple docstring'''




								pass



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





								'''simple docstring'''




								pass



				@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""      )
				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[Any]      )  ->    Tuple:





								'''simple docstring'''




								pass



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





								'''simple docstring'''




								for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												_lowerCAmelCase    :    str										=      ViTMAEModel.from_pretrained(_UpperCAmelCase      )
												self.assertIsNotNone(_UpperCAmelCase      )






def 							_UpperCAmelCase  ():
				'''simple docstring'''
				_lowerCAmelCase    :    List[Any]										=      Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""       )
				return image





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



				@cached_property
				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple      )  ->    int:





								'''simple docstring'''




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



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





								'''simple docstring'''




								np.random.seed(2      )

								_lowerCAmelCase    :    str										=      ViTMAEForPreTraining.from_pretrained("""facebook/vit-mae-base"""      ).to(_UpperCAmelCase      )

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

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

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

								# verify the logits
								_lowerCAmelCase    :    Optional[Any]										=      torch.Size((1, 196, 768)      )
								self.assertEqual(outputs.logits.shape		,	_UpperCAmelCase      )

								_lowerCAmelCase    :    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(_UpperCAmelCase      )		,	atol=1E-4      )      )

