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


SCREAMING_SNAKE_CASE :	Optional[int]							     =		logging.get_logger(__name__)

# General docstring
SCREAMING_SNAKE_CASE :	Optional[Any]							     =		"""PoolFormerConfig"""

# Base docstring
SCREAMING_SNAKE_CASE :	List[Any]							     =		"""sail/poolformer_s12"""
SCREAMING_SNAKE_CASE :	Tuple							     =		[1, 512, 7, 7]

# Image classification docstring
SCREAMING_SNAKE_CASE :	Union[str, Any]							     =		"""sail/poolformer_s12"""
SCREAMING_SNAKE_CASE :	List[str]							     =		"""tabby, tabby cat"""

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



def   __A		(			_A							,				_A = 0.0							,				_A = False					):



 """simple docstring"""


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

  def __init__(		self :      Optional[int]     ,							__SCREAMING_SNAKE_CASE :      Optional[float] = None       ):
   super().__init__()
   __a       =  drop_prob

  def 		_UpperCAmelCase   (		self :      List[Any]     ,							__SCREAMING_SNAKE_CASE :      torch.Tensor       ):
   return drop_path(__SCREAMING_SNAKE_CASE     ,							self.drop_prob     ,							self.training       )

  def 		_UpperCAmelCase   (		self :      List[Any]       ):
   return "p={}".format(self.drop_prob       )
class    A_ (       nn.Module ):

  def __init__(		self :      int     ,							__SCREAMING_SNAKE_CASE :      Tuple     ,							__SCREAMING_SNAKE_CASE :      List[str]     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      List[str]     ,							__SCREAMING_SNAKE_CASE :      Union[str, Any]     ,							__SCREAMING_SNAKE_CASE :      Optional[int]=None       ):
   super().__init__()
   __a       =  patch_size if isinstance(__SCREAMING_SNAKE_CASE     ,							collections.abc.Iterable       ) else (patch_size, patch_size)
   __a       =  stride if isinstance(__SCREAMING_SNAKE_CASE     ,							collections.abc.Iterable       ) else (stride, stride)
   __a       =  padding if isinstance(__SCREAMING_SNAKE_CASE     ,							collections.abc.Iterable       ) else (padding, padding)

   __a       =  nn.Convad(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							kernel_size=__SCREAMING_SNAKE_CASE     ,							stride=__SCREAMING_SNAKE_CASE     ,							padding=__SCREAMING_SNAKE_CASE       )
   __a       =  norm_layer(__SCREAMING_SNAKE_CASE       ) if norm_layer else nn.Identity()

  def 		_UpperCAmelCase   (		self :      Union[str, Any]     ,							__SCREAMING_SNAKE_CASE :      Any       ):
   __a       =  self.projection(__SCREAMING_SNAKE_CASE       )
   __a       =  self.norm(__SCREAMING_SNAKE_CASE       )
   return embeddings
class    A_ (       nn.GroupNorm ):

  def __init__(		self :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]     ,							**__SCREAMING_SNAKE_CASE :      Optional[Any]       ):
   super().__init__(1     ,							__SCREAMING_SNAKE_CASE     ,							**__SCREAMING_SNAKE_CASE       )
class    A_ (       nn.Module ):

  def __init__(		self :      List[str]     ,							__SCREAMING_SNAKE_CASE :      Dict       ):
   super().__init__()
   __a       =  nn.AvgPoolad(__SCREAMING_SNAKE_CASE     ,							stride=1     ,							padding=pool_size // 2     ,							count_include_pad=__SCREAMING_SNAKE_CASE       )

  def 		_UpperCAmelCase   (		self :      str     ,							__SCREAMING_SNAKE_CASE :      Optional[int]       ):
   return self.pool(__SCREAMING_SNAKE_CASE       ) - hidden_states
class    A_ (       nn.Module ):

  def __init__(		self :      Union[str, Any]     ,							__SCREAMING_SNAKE_CASE :      Tuple     ,							__SCREAMING_SNAKE_CASE :      List[str]     ,							__SCREAMING_SNAKE_CASE :      int     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]       ):
   super().__init__()
   __a       =  nn.Convad(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							1       )
   __a       =  nn.Convad(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							1       )
   __a       =  PoolFormerDropPath(__SCREAMING_SNAKE_CASE       )
   if isinstance(config.hidden_act     ,							__SCREAMING_SNAKE_CASE       ):
    __a       =  ACTaFN[config.hidden_act]
   else:
    __a       =  config.hidden_act

  def 		_UpperCAmelCase   (		self :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      str       ):
   __a       =  self.conva(__SCREAMING_SNAKE_CASE       )
   __a       =  self.act_fn(__SCREAMING_SNAKE_CASE       )
   __a       =  self.drop(__SCREAMING_SNAKE_CASE       )
   __a       =  self.conva(__SCREAMING_SNAKE_CASE       )
   __a       =  self.drop(__SCREAMING_SNAKE_CASE       )

   return hidden_states
class    A_ (       nn.Module ):

  def __init__(		self :      Tuple     ,							__SCREAMING_SNAKE_CASE :      List[str]     ,							__SCREAMING_SNAKE_CASE :      List[str]     ,							__SCREAMING_SNAKE_CASE :      Optional[int]     ,							__SCREAMING_SNAKE_CASE :      Dict     ,							__SCREAMING_SNAKE_CASE :      str     ,							__SCREAMING_SNAKE_CASE :      str       ):
   super().__init__()
   __a       =  PoolFormerPooling(__SCREAMING_SNAKE_CASE       )
   __a       =  PoolFormerOutput(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )
   __a       =  PoolFormerGroupNorm(__SCREAMING_SNAKE_CASE       )
   __a       =  PoolFormerGroupNorm(__SCREAMING_SNAKE_CASE       )

   # Useful for training neural nets
   __a       =  PoolFormerDropPath(__SCREAMING_SNAKE_CASE       ) if drop_path > 0.0 else nn.Identity()
   __a       =  config.use_layer_scale
   if config.use_layer_scale:
    __a       =  nn.Parameter(
        config.layer_scale_init_value * torch.ones((__SCREAMING_SNAKE_CASE)       )     ,							requires_grad=__SCREAMING_SNAKE_CASE       )
    __a       =  nn.Parameter(
        config.layer_scale_init_value * torch.ones((__SCREAMING_SNAKE_CASE)       )     ,							requires_grad=__SCREAMING_SNAKE_CASE       )

  def 		_UpperCAmelCase   (		self :      int     ,							__SCREAMING_SNAKE_CASE :      str       ):
   if self.use_layer_scale:
    __a       =  self.pooling(self.before_norm(__SCREAMING_SNAKE_CASE       )       )
    __a       =  self.layer_scale_a.unsqueeze(-1       ).unsqueeze(-1       ) * pooling_output
    # First residual connection
    __a       =  hidden_states + self.drop_path(__SCREAMING_SNAKE_CASE       )
    __a       =  ()

    __a       =  self.output(self.after_norm(__SCREAMING_SNAKE_CASE       )       )
    __a       =  self.layer_scale_a.unsqueeze(-1       ).unsqueeze(-1       ) * layer_output
    # Second residual connection
    __a       =  hidden_states + self.drop_path(__SCREAMING_SNAKE_CASE       )

    __a       =  (output,) + outputs
    return outputs

   else:
    __a       =  self.drop_path(self.pooling(self.before_norm(__SCREAMING_SNAKE_CASE       )       )       )
    # First residual connection
    __a       =  pooling_output + hidden_states
    __a       =  ()

    # Second residual connection inside the PoolFormerOutput block
    __a       =  self.drop_path(self.output(self.after_norm(__SCREAMING_SNAKE_CASE       )       )       )
    __a       =  hidden_states + layer_output

    __a       =  (output,) + outputs
    return outputs
class    A_ (       nn.Module ):

  def __init__(		self :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      Any       ):
   super().__init__()
   __a       =  config
   # stochastic depth decay rule
   __a       =  [x.item() for x in torch.linspace(0     ,							config.drop_path_rate     ,							sum(config.depths       )       )]

   # patch embeddings
   __a       =  []
   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       =  nn.ModuleList(__SCREAMING_SNAKE_CASE       )

   # Transformer blocks
   __a       =  []
   __a       =  0
   for i in range(config.num_encoder_blocks       ):
    # each block consists of layers
    __a       =  []
    if i != 0:
     cur += config.depths[i - 1]
    for j in range(config.depths[i]       ):
     layers.append(
         PoolFormerLayer(
             __SCREAMING_SNAKE_CASE     ,							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(__SCREAMING_SNAKE_CASE       )       )

   __a       =  nn.ModuleList(__SCREAMING_SNAKE_CASE       )

  def 		_UpperCAmelCase   (		self :      Any     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      Any=False     ,							__SCREAMING_SNAKE_CASE :      List[Any]=True       ):
   __a       =  () if output_hidden_states else None

   __a       =  pixel_values
   for idx, layers in enumerate(zip(self.patch_embeddings     ,							self.block       )       ):
    __a       ,     __a       =  layers
    # Get patch embeddings from hidden_states
    __a       =  embedding_layer(__SCREAMING_SNAKE_CASE       )
    # Send the embeddings through the blocks
    for _, blk in enumerate(__SCREAMING_SNAKE_CASE       ):
     __a       =  blk(__SCREAMING_SNAKE_CASE       )
     __a       =  layer_outputs[0]

    if output_hidden_states:
     __a       =  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=__SCREAMING_SNAKE_CASE     ,							hidden_states=__SCREAMING_SNAKE_CASE       )
class    A_ (       a_ ):
  _SCREAMING_SNAKE_CASE    				=			PoolFormerConfig
  _SCREAMING_SNAKE_CASE    				=			"""poolformer"""
  _SCREAMING_SNAKE_CASE    				=			"""pixel_values"""
  _SCREAMING_SNAKE_CASE    				=			True

  def 		_UpperCAmelCase   (		self :      int     ,							__SCREAMING_SNAKE_CASE :      Any       ):
   if isinstance(__SCREAMING_SNAKE_CASE     ,							(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(__SCREAMING_SNAKE_CASE     ,							nn.LayerNorm       ):
    module.bias.data.zero_()
    module.weight.data.fill_(1.0       )

  def 		_UpperCAmelCase   (		self :      Dict     ,							__SCREAMING_SNAKE_CASE :      List[Any]     ,							__SCREAMING_SNAKE_CASE :      int=False       ):
   if isinstance(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       ):
    __a       =  value
SCREAMING_SNAKE_CASE :	Any							     =		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.
"""

SCREAMING_SNAKE_CASE :	Any							     =		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."""			,       a_			,       )
class    A_ (       a_ ):

  def __init__(		self :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      int       ):
   super().__init__(__SCREAMING_SNAKE_CASE       )
   __a       =  config

   __a       =  PoolFormerEncoder(__SCREAMING_SNAKE_CASE       )

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

  def 		_UpperCAmelCase   (		self :      Tuple       ):
   return self.embeddings.patch_embeddings

  @add_start_docstrings_to_model_forward(__SCREAMING_SNAKE_CASE       )
  @add_code_sample_docstrings(
      checkpoint=_CHECKPOINT_FOR_DOC     ,							output_type=__SCREAMING_SNAKE_CASE     ,							config_class=_CONFIG_FOR_DOC     ,							modality="vision"     ,							expected_output=_EXPECTED_OUTPUT_SHAPE     ,							)
  def 		_UpperCAmelCase   (		self :      Tuple     ,							__SCREAMING_SNAKE_CASE :      Optional[torch.FloatTensor] = None     ,							__SCREAMING_SNAKE_CASE :      Optional[bool] = None     ,							__SCREAMING_SNAKE_CASE :      Optional[bool] = None     ,							):
   __a       =  (
       output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
   )
   __a       =  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       =  self.encoder(
       __SCREAMING_SNAKE_CASE     ,							output_hidden_states=__SCREAMING_SNAKE_CASE     ,							return_dict=__SCREAMING_SNAKE_CASE     ,							)
   __a       =  encoder_outputs[0]

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

   return BaseModelOutputWithNoAttention(
       last_hidden_state=__SCREAMING_SNAKE_CASE     ,							hidden_states=encoder_outputs.hidden_states     ,							)
class    A_ (       nn.Module ):

  def __init__(		self :      int     ,							__SCREAMING_SNAKE_CASE :      Any       ):
   super().__init__()
   __a       =  nn.Linear(config.hidden_size     ,							config.hidden_size       )

  def 		_UpperCAmelCase   (		self :      Dict     ,							__SCREAMING_SNAKE_CASE :      Optional[int]       ):
   __a       =  self.dense(__SCREAMING_SNAKE_CASE       )
   return output
@add_start_docstrings(
    """
    PoolFormer Model transformer with an image classification head on top
    """			,       a_			,       )
class    A_ (       a_ ):

  def __init__(		self :      str     ,							__SCREAMING_SNAKE_CASE :      int       ):
   super().__init__(__SCREAMING_SNAKE_CASE       )
   __a       =  config.num_labels
   __a       =  PoolFormerModel(__SCREAMING_SNAKE_CASE       )

   # Final norm
   __a       =  PoolFormerGroupNorm(config.hidden_sizes[-1]       )
   # Classifier head
   __a       =  (
       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(__SCREAMING_SNAKE_CASE       )
  @add_code_sample_docstrings(
      checkpoint=_IMAGE_CLASS_CHECKPOINT     ,							output_type=__SCREAMING_SNAKE_CASE     ,							config_class=_CONFIG_FOR_DOC     ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT     ,							)
  def 		_UpperCAmelCase   (		self :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      Optional[torch.FloatTensor] = None     ,							__SCREAMING_SNAKE_CASE :      Optional[torch.LongTensor] = None     ,							__SCREAMING_SNAKE_CASE :      Optional[bool] = None     ,							__SCREAMING_SNAKE_CASE :      Optional[bool] = None     ,							):
   __a       =  return_dict if return_dict is not None else self.config.use_return_dict

   __a       =  self.poolformer(
       __SCREAMING_SNAKE_CASE     ,							output_hidden_states=__SCREAMING_SNAKE_CASE     ,							return_dict=__SCREAMING_SNAKE_CASE     ,							)

   __a       =  outputs[0]

   __a       =  self.classifier(self.norm(__SCREAMING_SNAKE_CASE       ).mean([-2, -1]       )       )

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

    if self.config.problem_type == "regression":
     __a       =  MSELoss()
     if self.num_labels == 1:
      __a       =  loss_fct(logits.squeeze()     ,							labels.squeeze()       )
     else:
      __a       =  loss_fct(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )
    elif self.config.problem_type == "single_label_classification":
     __a       =  CrossEntropyLoss()
     __a       =  loss_fct(logits.view(-1     ,							self.num_labels       )     ,							labels.view(-1       )       )
    elif self.config.problem_type == "multi_label_classification":
     __a       =  BCEWithLogitsLoss()
     __a       =  loss_fct(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )

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

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



import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    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 ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
  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    A_ :

  def __init__(		self :      Union[str, Any]     ,							__SCREAMING_SNAKE_CASE :      Optional[int]     ,							__SCREAMING_SNAKE_CASE :      List[Any]=13     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]=30     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]=2     ,							__SCREAMING_SNAKE_CASE :      Optional[int]=3     ,							__SCREAMING_SNAKE_CASE :      List[Any]=True     ,							__SCREAMING_SNAKE_CASE :      Dict=True     ,							__SCREAMING_SNAKE_CASE :      Optional[int]=32     ,							__SCREAMING_SNAKE_CASE :      str=5     ,							__SCREAMING_SNAKE_CASE :      List[str]=4     ,							__SCREAMING_SNAKE_CASE :      List[Any]=37     ,							__SCREAMING_SNAKE_CASE :      int="gelu"     ,							__SCREAMING_SNAKE_CASE :      List[Any]=0.1     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]=0.1     ,							__SCREAMING_SNAKE_CASE :      List[Any]=10     ,							__SCREAMING_SNAKE_CASE :      int=0.02     ,							__SCREAMING_SNAKE_CASE :      Dict=None     ,							__SCREAMING_SNAKE_CASE :      str=2     ,							):
   __a       =  parent
   __a       =  batch_size
   __a       =  image_size
   __a       =  patch_size
   __a       =  num_channels
   __a       =  is_training
   __a       =  use_labels
   __a       =  hidden_size
   __a       =  num_hidden_layers
   __a       =  num_attention_heads
   __a       =  intermediate_size
   __a       =  hidden_act
   __a       =  hidden_dropout_prob
   __a       =  attention_probs_dropout_prob
   __a       =  type_sequence_label_size
   __a       =  initializer_range
   __a       =  scope
   __a       =  encoder_stride

   # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
   __a       =  (image_size // patch_size) ** 2
   __a       =  num_patches + 1

  def 		_UpperCAmelCase   (		self :      Any       ):
   __a       =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

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

   __a       =  self.get_config()

   return config, pixel_values, labels

  def 		_UpperCAmelCase   (		self :      str       ):
   return ViTConfig(
       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=__SCREAMING_SNAKE_CASE     ,							initializer_range=self.initializer_range     ,							encoder_stride=self.encoder_stride     ,							)

  def 		_UpperCAmelCase   (		self :      Dict     ,							__SCREAMING_SNAKE_CASE :      List[Any]     ,							__SCREAMING_SNAKE_CASE :      Dict     ,							__SCREAMING_SNAKE_CASE :      Any       ):
   __a       =  ViTModel(config=__SCREAMING_SNAKE_CASE       )
   model.to(__SCREAMING_SNAKE_CASE       )
   model.eval()
   __a       =  model(__SCREAMING_SNAKE_CASE       )
   self.parent.assertEqual(result.last_hidden_state.shape     ,							(self.batch_size, self.seq_length, self.hidden_size)       )

  def 		_UpperCAmelCase   (		self :      Union[str, Any]     ,							__SCREAMING_SNAKE_CASE :      Tuple     ,							__SCREAMING_SNAKE_CASE :      str     ,							__SCREAMING_SNAKE_CASE :      Dict       ):
   __a       =  ViTForMaskedImageModeling(config=__SCREAMING_SNAKE_CASE       )
   model.to(__SCREAMING_SNAKE_CASE       )
   model.eval()
   __a       =  model(__SCREAMING_SNAKE_CASE       )
   self.parent.assertEqual(
       result.reconstruction.shape     ,							(self.batch_size, self.num_channels, self.image_size, self.image_size)       )

   # test greyscale images
   __a       =  1
   __a       =  ViTForMaskedImageModeling(__SCREAMING_SNAKE_CASE       )
   model.to(__SCREAMING_SNAKE_CASE       )
   model.eval()

   __a       =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
   __a       =  model(__SCREAMING_SNAKE_CASE       )
   self.parent.assertEqual(result.reconstruction.shape     ,							(self.batch_size, 1, self.image_size, self.image_size)       )

  def 		_UpperCAmelCase   (		self :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      Tuple     ,							__SCREAMING_SNAKE_CASE :      Union[str, Any]     ,							__SCREAMING_SNAKE_CASE :      List[str]       ):
   __a       =  self.type_sequence_label_size
   __a       =  ViTForImageClassification(__SCREAMING_SNAKE_CASE       )
   model.to(__SCREAMING_SNAKE_CASE       )
   model.eval()
   __a       =  model(__SCREAMING_SNAKE_CASE     ,							labels=__SCREAMING_SNAKE_CASE       )
   self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.type_sequence_label_size)       )

   # test greyscale images
   __a       =  1
   __a       =  ViTForImageClassification(__SCREAMING_SNAKE_CASE       )
   model.to(__SCREAMING_SNAKE_CASE       )
   model.eval()

   __a       =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
   __a       =  model(__SCREAMING_SNAKE_CASE       )
   self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.type_sequence_label_size)       )

  def 		_UpperCAmelCase   (		self :      List[Any]       ):
   __a       =  self.prepare_config_and_inputs()
   (
       (
       __a
   )       ,     (
       __a
   )       ,     (
       __a
   )       ,     
   )       =  config_and_inputs
   __a       =  {"pixel_values": pixel_values}
   return config, inputs_dict
@require_torch
class    A_ (       a_			,       a_			,       unittest.TestCase ):
  _SCREAMING_SNAKE_CASE    				=			(
      (
          ViTModel,
          ViTForImageClassification,
          ViTForMaskedImageModeling,
      )
      if is_torch_available()
      else ()
  )
  _SCREAMING_SNAKE_CASE    				=			(
      {"""feature-extraction""": ViTModel, """image-classification""": ViTForImageClassification}
      if is_torch_available()
      else {}
  )
  _SCREAMING_SNAKE_CASE    				=			True

  _SCREAMING_SNAKE_CASE    				=			False
  _SCREAMING_SNAKE_CASE    				=			False
  _SCREAMING_SNAKE_CASE    				=			False

  def 		_UpperCAmelCase   (		self :      Optional[Any]       ):
   __a       =  ViTModelTester(self       )
   __a       =  ConfigTester(self     ,							config_class=__SCREAMING_SNAKE_CASE     ,							has_text_modality=__SCREAMING_SNAKE_CASE     ,							hidden_size=37       )

  def 		_UpperCAmelCase   (		self :      List[str]       ):
   self.config_tester.run_common_tests()

  @unittest.skip(reason="ViT does not use inputs_embeds"       )
  def 		_UpperCAmelCase   (		self :      Optional[Any]       ):
   pass

  def 		_UpperCAmelCase   (		self :      int       ):
   __a       ,     __a       =  self.model_tester.prepare_config_and_inputs_for_common()

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

  def 		_UpperCAmelCase   (		self :      Optional[int]       ):
   __a       ,     __a       =  self.model_tester.prepare_config_and_inputs_for_common()

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

    __a       =  ["pixel_values"]
    self.assertListEqual(arg_names[:1]     ,							__SCREAMING_SNAKE_CASE       )

  def 		_UpperCAmelCase   (		self :      List[str]       ):
   __a       =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE       )

  def 		_UpperCAmelCase   (		self :      Any       ):
   __a       =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_image_modeling(*__SCREAMING_SNAKE_CASE       )

  def 		_UpperCAmelCase   (		self :      Tuple       ):
   __a       =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*__SCREAMING_SNAKE_CASE       )

  @slow
  def 		_UpperCAmelCase   (		self :      Dict       ):
   for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    __a       =  ViTModel.from_pretrained(__SCREAMING_SNAKE_CASE       )
    self.assertIsNotNone(__SCREAMING_SNAKE_CASE       )



def   __A		(			):



 """simple docstring"""


 __a       =  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"					)
 return image






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

  @cached_property
  def 		_UpperCAmelCase   (		self :      List[str]       ):
   return ViTImageProcessor.from_pretrained("google/vit-base-patch16-224"       ) if is_vision_available() else None

  @slow
  def 		_UpperCAmelCase   (		self :      str       ):
   __a       =  ViTForImageClassification.from_pretrained("google/vit-base-patch16-224"       ).to(__SCREAMING_SNAKE_CASE       )

   __a       =  self.default_image_processor
   __a       =  prepare_img()
   __a       =  image_processor(images=__SCREAMING_SNAKE_CASE     ,							return_tensors="pt"       ).to(__SCREAMING_SNAKE_CASE       )

   # forward pass
   with torch.no_grad():
    __a       =  model(**__SCREAMING_SNAKE_CASE       )

   # verify the logits
   __a       =  torch.Size((1, 10_00)       )
   self.assertEqual(outputs.logits.shape     ,							__SCREAMING_SNAKE_CASE       )

   __a       =  torch.tensor([-0.27_44, 0.82_15, -0.08_36]       ).to(__SCREAMING_SNAKE_CASE       )

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

  @slow
  def 		_UpperCAmelCase   (		self :      Union[str, Any]       ):
   # ViT models have an `interpolate_pos_encoding` argument in their forward method,
   # allowing to interpolate the pre-trained position embeddings in order to use
   # the model on higher resolutions. The DINO model by Facebook AI leverages this
   # to visualize self-attention on higher resolution images.
   __a       =  ViTModel.from_pretrained("facebook/dino-vits8"       ).to(__SCREAMING_SNAKE_CASE       )

   __a       =  ViTImageProcessor.from_pretrained("facebook/dino-vits8"     ,							size=4_80       )
   __a       =  prepare_img()
   __a       =  image_processor(images=__SCREAMING_SNAKE_CASE     ,							return_tensors="pt"       )
   __a       =  inputs.pixel_values.to(__SCREAMING_SNAKE_CASE       )

   # forward pass
   with torch.no_grad():
    __a       =  model(__SCREAMING_SNAKE_CASE     ,							interpolate_pos_encoding=__SCREAMING_SNAKE_CASE       )

   # verify the logits
   __a       =  torch.Size((1, 36_01, 3_84)       )
   self.assertEqual(outputs.last_hidden_state.shape     ,							__SCREAMING_SNAKE_CASE       )

   __a       =  torch.tensor(
       [[4.23_40, 4.39_06, -6.66_92], [4.54_63, 1.89_28, -6.72_57], [4.44_29, 0.84_96, -5.85_85]]       ).to(__SCREAMING_SNAKE_CASE       )

   self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]     ,							__SCREAMING_SNAKE_CASE     ,							atol=1E-4       )       )

  @slow
  @require_accelerate
  @require_torch_gpu
  def 		_UpperCAmelCase   (		self :      Any       ):
   __a       =  ViTModel.from_pretrained("facebook/dino-vits8"     ,							torch_dtype=torch.floataa     ,							device_map="auto"       )
   __a       =  self.default_image_processor

   __a       =  prepare_img()
   __a       =  image_processor(images=__SCREAMING_SNAKE_CASE     ,							return_tensors="pt"       )
   __a       =  inputs.pixel_values.to(__SCREAMING_SNAKE_CASE       )

   # forward pass to make sure inference works in fp16
   with torch.no_grad():
    __a       =  model(__SCREAMING_SNAKE_CASE       )


