






'''simple docstring'''
import torch

from diffusers import CMStochasticIterativeScheduler

from .test_schedulers import SchedulerCommonTest






class    __SCREAMING_SNAKE_CASE					(	_lowerCAmelCase ):
 snake_case_     							=							(CMStochasticIterativeScheduler,)
 snake_case_     							=							10




 def    __magic_name__		(					self		:							Optional[Any]    ,					**__lowercase		:							Tuple					)		->  Dict:
       SCREAMING_SNAKE_CASE__   : Optional[Any]													={
           "num_train_timesteps": 2_01,
           "sigma_min": 0.002,
           "sigma_max": 80.0,
       }

       config.update(**__lowercase					)
       return config




 def    __magic_name__		(					self		:							List[Any]					)		->  Optional[Any]:
       SCREAMING_SNAKE_CASE__   : Optional[Any]													=10

       SCREAMING_SNAKE_CASE__   : Dict													=self.get_scheduler_config()
       SCREAMING_SNAKE_CASE__   : Any													=self.scheduler_classes[0](**__lowercase					)

       scheduler.set_timesteps(__lowercase					)

       SCREAMING_SNAKE_CASE__   : Any													=scheduler.timesteps[0]
       SCREAMING_SNAKE_CASE__   : str													=scheduler.timesteps[1]

       SCREAMING_SNAKE_CASE__   : Optional[Any]													=self.dummy_sample
       SCREAMING_SNAKE_CASE__   : str													=0.1 * sample

       SCREAMING_SNAKE_CASE__   : Tuple													=scheduler.step(__lowercase    ,					__lowercase    ,					__lowercase					).prev_sample
       SCREAMING_SNAKE_CASE__   : str													=scheduler.step(__lowercase    ,					__lowercase    ,					__lowercase					).prev_sample

       self.assertEqual(output_a.shape    ,					sample.shape					)
       self.assertEqual(output_a.shape    ,					output_a.shape					)




 def    __magic_name__		(					self		:							str					)		->  int:




       for timesteps in [10, 50, 1_00, 10_00]:
             self.check_over_configs(num_train_timesteps=__lowercase					)




 def    __magic_name__		(					self		:							Any					)		->  List[Any]:




       for clip_denoised in [True, False]:
             self.check_over_configs(clip_denoised=__lowercase					)




 def    __magic_name__		(					self		:							List[str]					)		->  Union[str, Any]:
       SCREAMING_SNAKE_CASE__   : Optional[int]													=self.scheduler_classes[0]
       SCREAMING_SNAKE_CASE__   : Any													=self.get_scheduler_config()
       SCREAMING_SNAKE_CASE__   : Optional[int]													=scheduler_class(**__lowercase					)

       SCREAMING_SNAKE_CASE__   : int													=1
       scheduler.set_timesteps(__lowercase					)
       SCREAMING_SNAKE_CASE__   : List[Any]													=scheduler.timesteps

       SCREAMING_SNAKE_CASE__   : List[Any]													=torch.manual_seed(0					)

       SCREAMING_SNAKE_CASE__   : int													=self.dummy_model()
       SCREAMING_SNAKE_CASE__   : Optional[int]													=self.dummy_sample_deter * scheduler.init_noise_sigma

       for i, t in enumerate(__lowercase					):
             # 1. scale model input
             SCREAMING_SNAKE_CASE__   : Tuple													=scheduler.scale_model_input(__lowercase    ,					__lowercase					)

             # 2. predict noise residual
             SCREAMING_SNAKE_CASE__   : Union[str, Any]													=model(__lowercase    ,					__lowercase					)

             # 3. predict previous sample x_t-1
             SCREAMING_SNAKE_CASE__   : List[Any]													=scheduler.step(__lowercase    ,					__lowercase    ,					__lowercase    ,					generator=__lowercase					).prev_sample

             SCREAMING_SNAKE_CASE__   : Any													=pred_prev_sample

       SCREAMING_SNAKE_CASE__   : str													=torch.sum(torch.abs(__lowercase					)					)
       SCREAMING_SNAKE_CASE__   : str													=torch.mean(torch.abs(__lowercase					)					)

       assert abs(result_sum.item() - 1_92.76_14					) < 1e-2
       assert abs(result_mean.item() - 0.2510					) < 1e-3




 def    __magic_name__		(					self		:							Optional[Any]					)		->  List[str]:
       SCREAMING_SNAKE_CASE__   : int													=self.scheduler_classes[0]
       SCREAMING_SNAKE_CASE__   : Dict													=self.get_scheduler_config()
       SCREAMING_SNAKE_CASE__   : Union[str, Any]													=scheduler_class(**__lowercase					)

       SCREAMING_SNAKE_CASE__   : List[Any]													=[1_06, 0]
       scheduler.set_timesteps(timesteps=__lowercase					)
       SCREAMING_SNAKE_CASE__   : List[Any]													=scheduler.timesteps

       SCREAMING_SNAKE_CASE__   : Dict													=torch.manual_seed(0					)

       SCREAMING_SNAKE_CASE__   : Dict													=self.dummy_model()
       SCREAMING_SNAKE_CASE__   : Tuple													=self.dummy_sample_deter * scheduler.init_noise_sigma

       for t in timesteps:
             # 1. scale model input
             SCREAMING_SNAKE_CASE__   : Optional[int]													=scheduler.scale_model_input(__lowercase    ,					__lowercase					)

             # 2. predict noise residual
             SCREAMING_SNAKE_CASE__   : int													=model(__lowercase    ,					__lowercase					)

             # 3. predict previous sample x_t-1
             SCREAMING_SNAKE_CASE__   : int													=scheduler.step(__lowercase    ,					__lowercase    ,					__lowercase    ,					generator=__lowercase					).prev_sample

             SCREAMING_SNAKE_CASE__   : str													=pred_prev_sample

       SCREAMING_SNAKE_CASE__   : Any													=torch.sum(torch.abs(__lowercase					)					)
       SCREAMING_SNAKE_CASE__   : str													=torch.mean(torch.abs(__lowercase					)					)

       assert abs(result_sum.item() - 3_47.63_57					) < 1e-2
       assert abs(result_mean.item() - 0.4527					) < 1e-3




 def    __magic_name__		(					self		:							Any					)		->  Tuple:
       SCREAMING_SNAKE_CASE__   : Any													=self.scheduler_classes[0]
       SCREAMING_SNAKE_CASE__   : Union[str, Any]													=self.get_scheduler_config()
       SCREAMING_SNAKE_CASE__   : Optional[Any]													=scheduler_class(**__lowercase					)

       SCREAMING_SNAKE_CASE__   : List[Any]													=[39, 30, 12, 15, 0]

       with self.assertRaises(__lowercase    ,					msg='''`timesteps` must be in descending order.'''					):
             scheduler.set_timesteps(timesteps=__lowercase					)




 def    __magic_name__		(					self		:							Dict					)		->  Any:
       SCREAMING_SNAKE_CASE__   : Dict													=self.scheduler_classes[0]
       SCREAMING_SNAKE_CASE__   : int													=self.get_scheduler_config()
       SCREAMING_SNAKE_CASE__   : Union[str, Any]													=scheduler_class(**__lowercase					)

       SCREAMING_SNAKE_CASE__   : str													=[39, 30, 12, 1, 0]
       SCREAMING_SNAKE_CASE__   : int													=len(__lowercase					)

       with self.assertRaises(__lowercase    ,					msg='''Can only pass one of `num_inference_steps` or `timesteps`.'''					):
             scheduler.set_timesteps(num_inference_steps=__lowercase    ,					timesteps=__lowercase					)




 def    __magic_name__		(					self		:							Any					)		->  Optional[Any]:
       SCREAMING_SNAKE_CASE__   : Dict													=self.scheduler_classes[0]
       SCREAMING_SNAKE_CASE__   : List[str]													=self.get_scheduler_config()
       SCREAMING_SNAKE_CASE__   : List[str]													=scheduler_class(**__lowercase					)

       SCREAMING_SNAKE_CASE__   : Optional[int]													=[scheduler.config.num_train_timesteps]

       with self.assertRaises(
           __lowercase    ,					msg='''`timesteps` must start before `self.config.train_timesteps`: {scheduler.config.num_train_timesteps}}'''    ,					):
             scheduler.set_timesteps(timesteps=__lowercase					)






'''simple docstring'''


import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VideoMAEConfig
from transformers.models.auto import get_values
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 (
     MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
     VideoMAEForPreTraining,
     VideoMAEForVideoClassification,
     VideoMAEModel,
 )
 from transformers.models.videomae.modeling_videomae import VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from transformers import VideoMAEImageProcessor
class       a_   :





     def __init__( self		:					Any	,      lowercase		:					Optional[int]	,      lowercase		:					List[Any]=13	,      lowercase		:					int=10	,      lowercase		:					str=3	,      lowercase		:					List[Any]=2	,      lowercase		:					Dict=2	,      lowercase		:					List[str]=2	,      lowercase		:					int=True	,      lowercase		:					List[Any]=True	,      lowercase		:					Union[str, Any]=32	,      lowercase		:					Optional[int]=5	,      lowercase		:					List[Any]=4	,      lowercase		:					List[str]=37	,      lowercase		:					Union[str, Any]="gelu"	,      lowercase		:					List[Any]=0.1	,      lowercase		:					Any=0.1	,      lowercase		:					Optional[Any]=10	,      lowercase		:					Union[str, Any]=0.02	,      lowercase		:					Optional[int]=0.9	,      lowercase		:					List[str]=None	,      ):





        """simple docstring"""




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

        # in VideoMAE, the number of tokens equals num_frames/tubelet_size * num_patches per frame
        lowercase_			:str    			= (image_size // patch_size) ** 2
        lowercase_			:Union[str, Any]    			= (num_frames // tubelet_size) * self.num_patches_per_frame

        # use this variable to define bool_masked_pos
        lowercase_			:Optional[Any]    			= int(mask_ratio * self.seq_length						)





     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= floats_tensor(
            [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]						)

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

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

        return config, pixel_values, labels





     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        return VideoMAEConfig(
            image_size=self.image_size	,      patch_size=self.patch_size	,      num_channels=self.num_channels	,      num_frames=self.num_frames	,      tubelet_size=self.tubelet_size	,      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=lowercase	,      initializer_range=self.initializer_range	,      )





     def 							lowercase__      ( self		:					Optional[int]	,      lowercase		:					Dict	,      lowercase		:					Dict	,      lowercase		:					Optional[int]						):





        """simple docstring"""




        lowercase_			:int    			= VideoMAEModel(config=lowercase						)
        model.to(lowercase						)
        model.eval()
        lowercase_			:Optional[int]    			= model(lowercase						)
        self.parent.assertEqual(result.last_hidden_state.shape	,      (self.batch_size, self.seq_length, self.hidden_size)						)





     def 							lowercase__      ( self		:					str	,      lowercase		:					str	,      lowercase		:					List[str]	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= VideoMAEForPreTraining(lowercase						)
        model.to(lowercase						)
        model.eval()
        # important: each video needs to have the same number of masked patches
        # hence we define a single mask, which we then repeat for each example in the batch
        lowercase_			:Optional[int]    			= torch.ones((self.num_masks,)						)
        lowercase_			:List[str]    			= torch.cat([mask, torch.zeros(self.seq_length - mask.size(0						)						)]						)
        lowercase_			:Dict    			= mask.expand(self.batch_size	,      -1						).bool()

        lowercase_			:str    			= model(lowercase	,      lowercase						)
        # model only returns predictions for masked patches
        lowercase_			:Any    			= mask.sum().item()
        lowercase_			:Tuple    			= 3 * self.tubelet_size * self.patch_size**2
        self.parent.assertEqual(result.logits.shape	,      (self.batch_size, num_masked_patches, decoder_num_labels)						)







     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        lowercase_			:int    			= self.prepare_config_and_inputs()
        lowercase_			,							lowercase_			,							lowercase_			:Dict    			= config_and_inputs
        lowercase_			:Any    			= {"pixel_values": pixel_values}
        return config, inputs_dict







@require_torch
class       a_   (				_lowerCAmelCase	,						_lowerCAmelCase	,						unittest.TestCase							):
     __A      =			(
         (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
     )
     __A      =			(
         {"feature-extraction": VideoMAEModel, "video-classification": VideoMAEForVideoClassification}
         if is_torch_available()
         else {}
     )

     __A      =			False
     __A      =			False
     __A      =			False
     __A      =			False





     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        lowercase_			:List[Any]    			= VideoMAEModelTester(self						)
        lowercase_			:Dict    			= ConfigTester(self	,      config_class=lowercase	,      has_text_modality=lowercase	,      hidden_size=37						)





     def 							lowercase__      ( self		:					List[Any]	,      lowercase		:					List[str]	,      lowercase		:					List[str]	,      lowercase		:					List[str]=False						):





        """simple docstring"""




        lowercase_			:Tuple    			= copy.deepcopy(lowercase						)

        if model_class == VideoMAEForPreTraining:
           # important: each video needs to have the same number of masked patches
           # hence we define a single mask, which we then repeat for each example in the batch
           lowercase_			:Tuple    			= torch.ones((self.model_tester.num_masks,)						)
           lowercase_			:Tuple    			= torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0						)						)]						)
           lowercase_			:Optional[int]    			= mask.expand(self.model_tester.batch_size	,      -1						).bool()
           lowercase_			:Dict    			= bool_masked_pos.to(lowercase						)

        if return_labels:
           if model_class in [
               *get_values(lowercase						),
           ]:
              lowercase_			:Union[str, Any]    			= torch.zeros(
                  self.model_tester.batch_size	,      dtype=torch.long	,      device=lowercase						)

        return inputs_dict





     def 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        self.config_tester.run_common_tests()





     @unittest.skip(reason="VideoMAE does not use inputs_embeds"						)
     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        pass





     def 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        lowercase_			,							lowercase_			:List[Any]    			= self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           lowercase_			:Dict    			= model_class(lowercase						)
           self.assertIsInstance(model.get_input_embeddings()	,      (nn.Module)						)
           lowercase_			:List[Any]    			= model.get_output_embeddings()
           self.assertTrue(x is None or isinstance(lowercase	,      nn.Linear						)						)





     def 							lowercase__      ( self		:					Optional[int]						):





        """simple docstring"""




        lowercase_			,							lowercase_			:Optional[int]    			= self.model_tester.prepare_config_and_inputs_for_common()

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

           lowercase_			:str    			= ["pixel_values"]
           self.assertListEqual(arg_names[:1]	,      lowercase						)





     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        lowercase_			:Tuple    			= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase						)





     def 							lowercase__      ( self		:					Optional[int]						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*lowercase						)





     @slow
     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowercase_			:List[Any]    			= VideoMAEModel.from_pretrained(lowercase						)
           self.assertIsNotNone(lowercase						)





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





        """simple docstring"""




        if not self.has_attentions:
           pass

        else:
           lowercase_			,							lowercase_			:Any    			= self.model_tester.prepare_config_and_inputs_for_common()
           lowercase_			:Union[str, Any]    			= True

           for model_class in self.all_model_classes:
              lowercase_			:Dict    			= self.model_tester.seq_length - self.model_tester.num_masks
              lowercase_			:Optional[Any]    			= (
                  num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
              )

              lowercase_			:Union[str, Any]    			= True
              lowercase_			:List[Any]    			= False
              lowercase_			:Optional[int]    			= True
              lowercase_			:Union[str, Any]    			= model_class(lowercase						)
              model.to(lowercase						)
              model.eval()
              with torch.no_grad():
                 lowercase_			:List[Any]    			= model(**self._prepare_for_class(lowercase	,      lowercase						)						)
              lowercase_			:str    			= outputs.attentions
              self.assertEqual(len(lowercase						)	,      self.model_tester.num_hidden_layers						)

              # check that output_attentions also work using config
              del inputs_dict["output_attentions"]
              lowercase_			:Union[str, Any]    			= True
              lowercase_			:Optional[Any]    			= model_class(lowercase						)
              model.to(lowercase						)
              model.eval()
              with torch.no_grad():
                 lowercase_			:Optional[Any]    			= model(**self._prepare_for_class(lowercase	,      lowercase						)						)
              lowercase_			:Union[str, Any]    			= outputs.attentions
              self.assertEqual(len(lowercase						)	,      self.model_tester.num_hidden_layers						)

              self.assertListEqual(
                  list(attentions[0].shape[-3:]						)	,      [self.model_tester.num_attention_heads, seq_len, seq_len]	,      )
              lowercase_			:List[str]    			= len(lowercase						)

              # Check attention is always last and order is fine
              lowercase_			:Optional[Any]    			= True
              lowercase_			:Dict    			= True
              lowercase_			:Dict    			= model_class(lowercase						)
              model.to(lowercase						)
              model.eval()
              with torch.no_grad():
                 lowercase_			:List[Any]    			= model(**self._prepare_for_class(lowercase	,      lowercase						)						)

              self.assertEqual(out_len + 1	,      len(lowercase						)						)

              lowercase_			:int    			= outputs.attentions

              self.assertEqual(len(lowercase						)	,      self.model_tester.num_hidden_layers						)
              self.assertListEqual(
                  list(self_attentions[0].shape[-3:]						)	,      [self.model_tester.num_attention_heads, seq_len, seq_len]	,      )





     def 							lowercase__      ( self		:					int						):





        """simple docstring"""




        def check_hidden_states_output(lowercase		:					Union[str, Any]	,      lowercase		:					Dict	,      lowercase		:					Any						):
           lowercase_			:Any    			= model_class(lowercase						)
           model.to(lowercase						)
           model.eval()

           with torch.no_grad():
              lowercase_			:Optional[int]    			= model(**self._prepare_for_class(lowercase	,      lowercase						)						)

           lowercase_			:Optional[int]    			= outputs.hidden_states
           lowercase_			:Any    			= self.model_tester.num_hidden_layers + 1
           self.assertEqual(len(lowercase						)	,      lowercase						)

           lowercase_			:List[str]    			= self.model_tester.seq_length - self.model_tester.num_masks
           lowercase_			:List[Any]    			= num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length

           self.assertListEqual(
               list(hidden_states[0].shape[-2:]						)	,      [seq_length, self.model_tester.hidden_size]	,      )

        lowercase_			,							lowercase_			:List[Any]    			= self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
           lowercase_			:Optional[int]    			= True
           check_hidden_states_output(lowercase	,      lowercase	,      lowercase						)

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           lowercase_			:List[Any]    			= True

           check_hidden_states_output(lowercase	,      lowercase	,      lowercase						)







     @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."						)
     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        pass







def 							UpperCAmelCase_  (						):
   lowercase_			:Dict    			= hf_hub_download(
       repo_id="hf-internal-testing/spaghetti-video"	,filename="eating_spaghetti.npy"	,repo_type="dataset" )
   lowercase_			:Optional[Any]    			= np.load(__lowerCamelCase )
   return list(__lowerCamelCase )


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





     @cached_property
     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        return (
            VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]	,      image_std=[0.5, 0.5, 0.5]						)
            if is_vision_available()
            else None
        )





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





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= VideoMAEForVideoClassification.from_pretrained("MCG-NJU/videomae-base-finetuned-kinetics"						).to(
            lowercase						)

        lowercase_			:List[str]    			= self.default_image_processor
        lowercase_			:List[str]    			= prepare_video()
        lowercase_			:int    			= image_processor(lowercase	,      return_tensors="pt"						).to(lowercase						)

        # forward pass
        with torch.no_grad():
           lowercase_			:Dict    			= model(**lowercase						)

        # verify the logits
        lowercase_			:Dict    			= torch.Size((1, 400)						)
        self.assertEqual(outputs.logits.shape	,      lowercase						)

        lowercase_			:int    			= torch.tensor([0.36_69, -0.06_88, -0.24_21]						).to(lowercase						)

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







     @slow
     def 							lowercase__      ( self		:					str						):





        """simple docstring"""




        lowercase_			:List[Any]    			= VideoMAEForPreTraining.from_pretrained("MCG-NJU/videomae-base-short"						).to(lowercase						)

        lowercase_			:Dict    			= self.default_image_processor
        lowercase_			:Union[str, Any]    			= prepare_video()
        lowercase_			:List[str]    			= image_processor(lowercase	,      return_tensors="pt"						).to(lowercase						)

        # add boolean mask, indicating which patches to mask
        lowercase_			:int    			= hf_hub_download(repo_id="hf-internal-testing/bool-masked-pos"	,      filename="bool_masked_pos.pt"						)
        lowercase_			:List[str]    			= torch.load(lowercase						)

        # forward pass
        with torch.no_grad():
           lowercase_			:List[Any]    			= model(**lowercase						)

        # verify the logits
        lowercase_			:Union[str, Any]    			= torch.Size([1, 1_408, 1_536]						)
        lowercase_			:List[Any]    			= torch.tensor(
            [[0.79_94, 0.96_12, 0.85_08], [0.74_01, 0.89_58, 0.83_02], [0.58_62, 0.74_68, 0.73_25]]	,      device=lowercase						)
        self.assertEqual(outputs.logits.shape	,      lowercase						)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]	,      lowercase	,      atol=1e-4						)						)

        # verify the loss (`config.norm_pix_loss` = `True`)
        lowercase_			:Any    			= torch.tensor([0.51_42]	,      device=lowercase						)
        self.assertTrue(torch.allclose(outputs.loss	,      lowercase	,      atol=1e-4						)						)

        # verify the loss (`config.norm_pix_loss` = `False`)
        lowercase_			:Union[str, Any]    			= VideoMAEForPreTraining.from_pretrained("MCG-NJU/videomae-base-short"	,      norm_pix_loss=lowercase						).to(
            lowercase						)

        with torch.no_grad():
           lowercase_			:Tuple    			= model(**lowercase						)

        lowercase_			:Optional[Any]    			= torch.tensor(torch.tensor([0.64_69]						)	,      device=lowercase						)
        self.assertTrue(torch.allclose(outputs.loss	,      lowercase	,      atol=1e-4						)						)


