

"""simple docstring"""




import contextlib
import copy
import random
from typing import Any, Dict, Iterable, Optional, Union

import numpy as np
import torch

from .utils import deprecate, is_transformers_available


if is_transformers_available():
      import transformers
def   lowerCAmelCase						(					__UpperCamelCase      ):
 '''simple docstring'''





 random.seed(__UpperCamelCase      )
 np.random.seed(__UpperCamelCase      )
 torch.manual_seed(__UpperCamelCase      )
 torch.cuda.manual_seed_all(__UpperCamelCase      )
 # ^^ safe to call this function even if cuda is not available







class      __lowercase			:

  def __init__(  self		:			int					,A		:			Iterable[torch.nn.Parameter]					,A		:			float = 0.9_9_9_9					,A		:			float = 0.0					,A		:			int = 0					,A		:			bool = False					,A		:			Union[float, int] = 1.0					,A		:			Union[float, int] = 2 / 3					,A		:			Optional[Any] = None					,A		:			Dict[str, Any] = None					,**A		:			int					,):






   '''simple docstring'''

   if isinstance(A					,torch.nn.Module    ):
    UpperCAmelCase__					:  List[str]					   =					(
        """Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. """
        """Please pass the parameters of the module instead."""
    )
    deprecate(
        """passing a `torch.nn.Module` to `ExponentialMovingAverage`"""					,"""1.0.0"""					,A					,standard_warn=A					,)
    UpperCAmelCase__					:  Optional[int]					   =					parameters.parameters()

    # set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility
    UpperCAmelCase__					:  Any					   =					True

   if kwargs.get("""max_value"""					,A    ) is not None:
    UpperCAmelCase__					:  Optional[int]					   =					"""The `max_value` argument is deprecated. Please use `decay` instead."""
    deprecate("""max_value"""					,"""1.0.0"""					,A					,standard_warn=A    )
    UpperCAmelCase__					:  Tuple					   =					kwargs["""max_value"""]

   if kwargs.get("""min_value"""					,A    ) is not None:
    UpperCAmelCase__					:  List[str]					   =					"""The `min_value` argument is deprecated. Please use `min_decay` instead."""
    deprecate("""min_value"""					,"""1.0.0"""					,A					,standard_warn=A    )
    UpperCAmelCase__					:  List[Any]					   =					kwargs["""min_value"""]

   UpperCAmelCase__					:  int					   =					list(A    )
   UpperCAmelCase__					:  int					   =					[p.clone().detach() for p in parameters]

   if kwargs.get("""device"""					,A    ) is not None:
    UpperCAmelCase__					:  Union[str, Any]					   =					"""The `device` argument is deprecated. Please use `to` instead."""
    deprecate("""device"""					,"""1.0.0"""					,A					,standard_warn=A    )
    self.to(device=kwargs["""device"""]    )

   UpperCAmelCase__					:  Dict					   =					None

   UpperCAmelCase__					:  Dict					   =					decay
   UpperCAmelCase__					:  Any					   =					min_decay
   UpperCAmelCase__					:  int					   =					update_after_step
   UpperCAmelCase__					:  Optional[int]					   =					use_ema_warmup
   UpperCAmelCase__					:  str					   =					inv_gamma
   UpperCAmelCase__					:  List[Any]					   =					power
   UpperCAmelCase__					:  Tuple					   =					0
   UpperCAmelCase__					:  Union[str, Any]					   =					None  # set in `step()`

   UpperCAmelCase__					:  Tuple					   =					model_cls
   UpperCAmelCase__					:  Optional[int]					   =					model_config

  @classmethod
  def    __lowercase       (  cls		:			Dict					,A		:			Union[str, Any]					,A		:			str    ):






   '''simple docstring'''
   UpperCAmelCase__   ,      UpperCAmelCase__					:  Dict					   =					model_cls.load_config(A					,return_unused_kwargs=A    )
   UpperCAmelCase__					:  List[str]					   =					model_cls.from_pretrained(A    )

   UpperCAmelCase__					:  Dict					   =					cls(model.parameters()					,model_cls=A					,model_config=model.config    )

   ema_model.load_state_dict(A    )
   return ema_model

  def    __lowercase       (  self		:			List[Any]					,A		:			List[Any]    ):






   '''simple docstring'''
   if self.model_cls is None:
    raise ValueError("""`save_pretrained` can only be used if `model_cls` was defined at __init__."""    )

   if self.model_config is None:
    raise ValueError("""`save_pretrained` can only be used if `model_config` was defined at __init__."""    )

   UpperCAmelCase__					:  Any					   =					self.model_cls.from_config(self.model_config    )
   UpperCAmelCase__					:  Tuple					   =					self.state_dict()
   state_dict.pop("""shadow_params"""					,A    )

   model.register_to_config(**A    )
   self.copy_to(model.parameters()    )
   model.save_pretrained(A    )

  def    __lowercase       (  self		:			List[str]					,A		:			int    ):






   '''simple docstring'''
   UpperCAmelCase__					:  int					   =					max(0					,optimization_step - self.update_after_step - 1    )

   if step <= 0:
    return 0.0

   if self.use_ema_warmup:
    UpperCAmelCase__					:  Any					   =					1 - (1 + step / self.inv_gamma) ** -self.power
   else:
    UpperCAmelCase__					:  List[Any]					   =					(1 + step) / (10 + step)

   UpperCAmelCase__					:  Optional[Any]					   =					min(A					,self.decay    )
   # make sure decay is not smaller than min_decay
   UpperCAmelCase__					:  List[Any]					   =					max(A					,self.min_decay    )
   return cur_decay_value

  @torch.no_grad()
  def    __lowercase       (  self		:			Tuple					,A		:			Iterable[torch.nn.Parameter]    ):






   '''simple docstring'''
   if isinstance(A					,torch.nn.Module    ):
    UpperCAmelCase__					:  List[Any]					   =					(
        """Passing a `torch.nn.Module` to `ExponentialMovingAverage.step` is deprecated. """
        """Please pass the parameters of the module instead."""
    )
    deprecate(
        """passing a `torch.nn.Module` to `ExponentialMovingAverage.step`"""					,"""1.0.0"""					,A					,standard_warn=A					,)
    UpperCAmelCase__					:  List[str]					   =					parameters.parameters()

   UpperCAmelCase__					:  str					   =					list(A    )

   self.optimization_step += 1

   # Compute the decay factor for the exponential moving average.
   UpperCAmelCase__					:  str					   =					self.get_decay(self.optimization_step    )
   UpperCAmelCase__					:  Dict					   =					decay
   UpperCAmelCase__					:  List[str]					   =					1 - decay

   UpperCAmelCase__					:  Union[str, Any]					   =					contextlib.nullcontext
   if is_transformers_available() and transformers.deepspeed.is_deepspeed_zeroa_enabled():
    import deepspeed

   for s_param, param in zip(self.shadow_params					,A    ):
    if is_transformers_available() and transformers.deepspeed.is_deepspeed_zeroa_enabled():
     UpperCAmelCase__					:  Optional[Any]					   =					deepspeed.zero.GatheredParameters(A					,modifier_rank=A    )

    with context_manager():
     if param.requires_grad:
      s_param.sub_(one_minus_decay * (s_param - param)    )
     else:
      s_param.copy_(A    )

  def    __lowercase       (  self		:			Any					,A		:			Iterable[torch.nn.Parameter]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Any					   =					list(A    )
   for s_param, param in zip(self.shadow_params					,A    ):
    param.data.copy_(s_param.to(param.device    ).data    )

  def    __lowercase       (  self		:			Union[str, Any]					,A		:			Optional[Any]=None					,A		:			Dict=None    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					[
       p.to(device=A					,dtype=A    ) if p.is_floating_point() else p.to(device=A    )
       for p in self.shadow_params
   ]

  def    __lowercase       (  self		:			Dict    ):






   '''simple docstring'''
   return {
       "decay": self.decay,
       "min_decay": self.min_decay,
       "optimization_step": self.optimization_step,
       "update_after_step": self.update_after_step,
       "use_ema_warmup": self.use_ema_warmup,
       "inv_gamma": self.inv_gamma,
       "power": self.power,
       "shadow_params": self.shadow_params,
   }

  def    __lowercase       (  self		:			Optional[int]					,A		:			Iterable[torch.nn.Parameter]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[str]					   =					[param.detach().cpu().clone() for param in parameters]

  def    __lowercase       (  self		:			Dict					,A		:			Iterable[torch.nn.Parameter]    ):






   '''simple docstring'''
   if self.temp_stored_params is None:
    raise RuntimeError("""This ExponentialMovingAverage has no `store()`ed weights """ """to `restore()`"""    )
   for c_param, param in zip(self.temp_stored_params					,A    ):
    param.data.copy_(c_param.data    )

   # Better memory-wise.
   UpperCAmelCase__					:  Optional[int]					   =					None




  def    __lowercase       (  self		:			Optional[Any]					,A		:			dict    ):






   '''simple docstring'''
   UpperCAmelCase__					:  str					   =					copy.deepcopy(A    )

   UpperCAmelCase__					:  List[Any]					   =					state_dict.get("""decay"""					,self.decay    )
   if self.decay < 0.0 or self.decay > 1.0:
    raise ValueError("""Decay must be between 0 and 1"""    )

   UpperCAmelCase__					:  str					   =					state_dict.get("""min_decay"""					,self.min_decay    )
   if not isinstance(self.min_decay					,A    ):
    raise ValueError("""Invalid min_decay"""    )

   UpperCAmelCase__					:  str					   =					state_dict.get("""optimization_step"""					,self.optimization_step    )
   if not isinstance(self.optimization_step					,A    ):
    raise ValueError("""Invalid optimization_step"""    )

   UpperCAmelCase__					:  List[str]					   =					state_dict.get("""update_after_step"""					,self.update_after_step    )
   if not isinstance(self.update_after_step					,A    ):
    raise ValueError("""Invalid update_after_step"""    )

   UpperCAmelCase__					:  List[str]					   =					state_dict.get("""use_ema_warmup"""					,self.use_ema_warmup    )
   if not isinstance(self.use_ema_warmup					,A    ):
    raise ValueError("""Invalid use_ema_warmup"""    )

   UpperCAmelCase__					:  Optional[int]					   =					state_dict.get("""inv_gamma"""					,self.inv_gamma    )
   if not isinstance(self.inv_gamma					,(float, int)    ):
    raise ValueError("""Invalid inv_gamma"""    )

   UpperCAmelCase__					:  Any					   =					state_dict.get("""power"""					,self.power    )
   if not isinstance(self.power					,(float, int)    ):
    raise ValueError("""Invalid power"""    )

   UpperCAmelCase__					:  List[Any]					   =					state_dict.get("""shadow_params"""					,A    )
   if shadow_params is not None:
    UpperCAmelCase__					:  Union[str, Any]					   =					shadow_params
    if not isinstance(self.shadow_params					,A    ):
     raise ValueError("""shadow_params must be a list"""    )
    if not all(isinstance(A					,torch.Tensor    ) for p in self.shadow_params    ):
     raise ValueError("""shadow_params must all be Tensors"""    )








"""simple docstring"""




from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers import (
          TFDeiTForImageClassification,
          TFDeiTForImageClassificationWithTeacher,
          TFDeiTForMaskedImageModeling,
          TFDeiTModel,
      )
      from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import DeiTImageProcessor







class      __lowercase			:

  def __init__(  self		:			Optional[int]					,A		:			Optional[Any]					,A		:			int=13					,A		:			Any=30					,A		:			str=2					,A		:			Tuple=3					,A		:			Optional[Any]=True					,A		:			Union[str, Any]=True					,A		:			Dict=32					,A		:			Optional[int]=2					,A		:			Any=4					,A		:			str=37					,A		:			int="gelu"					,A		:			Dict=0.1					,A		:			str=0.1					,A		:			Optional[Any]=10					,A		:			List[str]=0.0_2					,A		:			Any=3					,A		:			Any=None					,A		:			int=2					,):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					parent
   UpperCAmelCase__					:  str					   =					batch_size
   UpperCAmelCase__					:  Tuple					   =					image_size
   UpperCAmelCase__					:  int					   =					patch_size
   UpperCAmelCase__					:  Dict					   =					num_channels
   UpperCAmelCase__					:  List[Any]					   =					is_training
   UpperCAmelCase__					:  Optional[Any]					   =					use_labels
   UpperCAmelCase__					:  Tuple					   =					hidden_size
   UpperCAmelCase__					:  List[str]					   =					num_hidden_layers
   UpperCAmelCase__					:  Dict					   =					num_attention_heads
   UpperCAmelCase__					:  List[Any]					   =					intermediate_size
   UpperCAmelCase__					:  Optional[Any]					   =					hidden_act
   UpperCAmelCase__					:  Optional[Any]					   =					hidden_dropout_prob
   UpperCAmelCase__					:  str					   =					attention_probs_dropout_prob
   UpperCAmelCase__					:  Dict					   =					type_sequence_label_size
   UpperCAmelCase__					:  Optional[int]					   =					initializer_range
   UpperCAmelCase__					:  Tuple					   =					scope
   UpperCAmelCase__					:  str					   =					encoder_stride

   # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
   UpperCAmelCase__					:  str					   =					(image_size // patch_size) ** 2
   UpperCAmelCase__					:  int					   =					num_patches + 2

  def    __lowercase       (  self		:			List[Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Union[str, Any]					   =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

   UpperCAmelCase__					:  List[str]					   =					None
   if self.use_labels:
    UpperCAmelCase__					:  Optional[Any]					   =					ids_tensor([self.batch_size]					,self.type_sequence_label_size    )

   UpperCAmelCase__					:  int					   =					self.get_config()

   return config, pixel_values, labels

  def    __lowercase       (  self		:			Optional[Any]    ):






   '''simple docstring'''
   return DeiTConfig(
       image_size=self.image_size					,patch_size=self.patch_size					,num_channels=self.num_channels					,hidden_size=self.hidden_size					,num_hidden_layers=self.num_hidden_layers					,num_attention_heads=self.num_attention_heads					,intermediate_size=self.intermediate_size					,hidden_act=self.hidden_act					,hidden_dropout_prob=self.hidden_dropout_prob					,attention_probs_dropout_prob=self.attention_probs_dropout_prob					,is_decoder=A					,initializer_range=self.initializer_range					,encoder_stride=self.encoder_stride					,)

  def    __lowercase       (  self		:			Union[str, Any]					,A		:			Optional[int]					,A		:			Optional[Any]					,A		:			Optional[Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Dict					   =					TFDeiTModel(config=A    )
   UpperCAmelCase__					:  Dict					   =					model(A    )
   self.parent.assertEqual(result.last_hidden_state.shape					,(self.batch_size, self.seq_length, self.hidden_size)    )

  def    __lowercase       (  self		:			Optional[int]					,A		:			Union[str, Any]					,A		:			Optional[int]					,A		:			str    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Dict					   =					TFDeiTForMaskedImageModeling(config=A    )
   UpperCAmelCase__					:  Any					   =					model(A    )
   self.parent.assertEqual(
       result.reconstruction.shape					,(self.batch_size, self.num_channels, self.image_size, self.image_size)    )

   # test greyscale images
   UpperCAmelCase__					:  Union[str, Any]					   =					1
   UpperCAmelCase__					:  Any					   =					TFDeiTForMaskedImageModeling(A    )

   UpperCAmelCase__					:  str					   =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
   UpperCAmelCase__					:  str					   =					model(A    )
   self.parent.assertEqual(result.reconstruction.shape					,(self.batch_size, 1, self.image_size, self.image_size)    )

  def    __lowercase       (  self		:			Dict					,A		:			Any					,A		:			Optional[Any]					,A		:			List[Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  int					   =					self.type_sequence_label_size
   UpperCAmelCase__					:  int					   =					TFDeiTForImageClassification(A    )
   UpperCAmelCase__					:  int					   =					model(A					,labels=A    )
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.type_sequence_label_size)    )

   # test greyscale images
   UpperCAmelCase__					:  Optional[int]					   =					1
   UpperCAmelCase__					:  int					   =					TFDeiTForImageClassification(A    )

   UpperCAmelCase__					:  List[str]					   =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
   UpperCAmelCase__					:  str					   =					model(A					,labels=A    )
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.type_sequence_label_size)    )




  def    __lowercase       (  self		:			int    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					self.prepare_config_and_inputs()
   UpperCAmelCase__   ,      UpperCAmelCase__   ,      UpperCAmelCase__					:  Optional[Any]					   =					config_and_inputs
   UpperCAmelCase__					:  Dict					   =					{"""pixel_values""": pixel_values}
   return config, inputs_dict







@require_tf
class      __lowercase			(    __lowerCamelCase	,			__lowerCamelCase	,			unittest.TestCase					):
  snake_case_				       =       (
      (
          TFDeiTModel,
          TFDeiTForImageClassification,
          TFDeiTForImageClassificationWithTeacher,
          TFDeiTForMaskedImageModeling,
      )
      if is_tf_available()
      else ()
  )
  snake_case_				       =       (
      {
          """feature-extraction""": TFDeiTModel,
          """image-classification""": (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
      }
      if is_tf_available()
      else {}
  )

  snake_case_				       =       False
  snake_case_				       =       False
  snake_case_				       =       False
  snake_case_				       =       False

  def    __lowercase       (  self		:			Optional[int]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					TFDeiTModelTester(self    )
   UpperCAmelCase__					:  int					   =					ConfigTester(self					,config_class=A					,has_text_modality=A					,hidden_size=37    )

  def    __lowercase       (  self		:			Union[str, Any]    ):






   '''simple docstring'''
   self.config_tester.run_common_tests()

  @unittest.skip(reason="""DeiT does not use inputs_embeds"""    )
  def    __lowercase       (  self		:			str    ):






   '''simple docstring'''
   pass

  def    __lowercase       (  self		:			Any    ):






   '''simple docstring'''
   UpperCAmelCase__   ,      UpperCAmelCase__					:  Dict					   =					self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    UpperCAmelCase__					:  int					   =					model_class(A    )
    self.assertIsInstance(model.get_input_embeddings()					,(tf.keras.layers.Layer)    )
    UpperCAmelCase__					:  Tuple					   =					model.get_output_embeddings()
    self.assertTrue(x is None or isinstance(A					,tf.keras.layers.Dense    )    )

  def    __lowercase       (  self		:			Tuple    ):






   '''simple docstring'''
   UpperCAmelCase__   ,      UpperCAmelCase__					:  List[Any]					   =					self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    UpperCAmelCase__					:  Dict					   =					model_class(A    )
    UpperCAmelCase__					:  Any					   =					inspect.signature(model.call    )
    # signature.parameters is an OrderedDict => so arg_names order is deterministic
    UpperCAmelCase__					:  List[Any]					   =					[*signature.parameters.keys()]

    UpperCAmelCase__					:  Optional[int]					   =					["""pixel_values"""]
    self.assertListEqual(arg_names[:1]					,A    )

  def    __lowercase       (  self		:			Any    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*A    )

  def    __lowercase       (  self		:			int    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Union[str, Any]					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_image_modeling(*A    )

  def    __lowercase       (  self		:			List[str]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  int					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*A    )

  def    __lowercase       (  self		:			str					,A		:			List[str]					,A		:			Optional[int]					,A		:			Union[str, Any]=False    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Any					   =					super()._prepare_for_class(A					,A					,return_labels=A    )

   if return_labels:
    if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call    ).parameters:
     del inputs_dict["labels"]

   return inputs_dict




  @slow
  def    __lowercase       (  self		:			List[str]    ):






   '''simple docstring'''
   for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    UpperCAmelCase__					:  List[Any]					   =					TFDeiTModel.from_pretrained(A    )
    self.assertIsNotNone(A    )







def   lowerCAmelCase						(					):
 '''simple docstring'''





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







@require_tf
@require_vision
class      __lowercase			(    unittest.TestCase					):

  @cached_property
  def    __lowercase       (  self		:			int    ):






   '''simple docstring'''
   return (
       DeiTImageProcessor.from_pretrained("""facebook/deit-base-distilled-patch16-224"""    )
       if is_vision_available()
       else None
   )




  @slow
  def    __lowercase       (  self		:			Optional[int]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  str					   =					TFDeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224"""    )

   UpperCAmelCase__					:  Union[str, Any]					   =					self.default_image_processor
   UpperCAmelCase__					:  Optional[int]					   =					prepare_img()
   UpperCAmelCase__					:  Optional[int]					   =					image_processor(images=A					,return_tensors="""tf"""    )

   # forward pass
   UpperCAmelCase__					:  Optional[Any]					   =					model(**A    )

   # verify the logits
   UpperCAmelCase__					:  int					   =					tf.TensorShape((1, 1_000)    )
   self.assertEqual(outputs.logits.shape					,A    )

   UpperCAmelCase__					:  Any					   =					tf.constant([-1.0_2_6_6, 0.1_9_1_2, -1.2_8_6_1]    )

   self.assertTrue(np.allclose(outputs.logits[0, :3]					,A					,atol=1e-4    )    )





