






from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)






@flax.struct.dataclass
class 		lowercase_					(	lowercase   ):






      '''simple docstring'''




      __snake_case       =  42
      __snake_case       =  42






class 		lowercase_					(	nn.Module   ):






      '''simple docstring'''




      __snake_case       =  42
      __snake_case       =  (16, 32, 96, 2_56)
      __snake_case       =  jnp.floataa







      def       __lowerCAmelCase					(      self      :						Dict  )    ->int:




       """simple docstring"""


       a		  =      nn.Conv(
           self.block_out_channels[0] ,					kernel_size=(3, 3) ,					padding=((1, 1), (1, 1)) ,					dtype=self.dtype ,					)

       a		  =      []
       for i in range(len(self.block_out_channels  ) - 1  ):
        a		  =      self.block_out_channels[i]
        a		  =      self.block_out_channels[i + 1]
        a		  =      nn.Conv(
            __UpperCAmelCase ,					kernel_size=(3, 3) ,					padding=((1, 1), (1, 1)) ,					dtype=self.dtype ,					)
        blocks.append(__UpperCAmelCase  )
        a		  =      nn.Conv(
            __UpperCAmelCase ,					kernel_size=(3, 3) ,					strides=(2, 2) ,					padding=((1, 1), (1, 1)) ,					dtype=self.dtype ,					)
        blocks.append(__UpperCAmelCase  )
       a		  =      blocks

       a		  =      nn.Conv(
           self.conditioning_embedding_channels ,					kernel_size=(3, 3) ,					padding=((1, 1), (1, 1)) ,					kernel_init=nn.initializers.zeros_init() ,					bias_init=nn.initializers.zeros_init() ,					dtype=self.dtype ,					)







      def __call__(      self      :						List[Any] ,					__UpperCAmelCase      :						Any  )    ->Dict:




       """simple docstring"""


       a		  =      self.conv_in(__UpperCAmelCase  )
       a		  =      nn.silu(__UpperCAmelCase  )

       for block in self.blocks:
        a		  =      block(__UpperCAmelCase  )
        a		  =      nn.silu(__UpperCAmelCase  )

       a		  =      self.conv_out(__UpperCAmelCase  )

       return embedding






@flax_register_to_config
class 		lowercase_					(	nn.Module ,  lowercase ,  lowercase   ):






      '''simple docstring'''




      __snake_case       =  32
      __snake_case       =  4
      __snake_case       =  (
          "CrossAttnDownBlock2D",
          "CrossAttnDownBlock2D",
          "CrossAttnDownBlock2D",
          "DownBlock2D",
      )
      __snake_case       =  False
      __snake_case       =  (3_20, 6_40, 12_80, 12_80)
      __snake_case       =  2
      __snake_case       =  8
      __snake_case       =  None
      __snake_case       =  12_80
      __snake_case       =  0.0
      __snake_case       =  False
      __snake_case       =  jnp.floataa
      __snake_case       =  True
      __snake_case       =  0
      __snake_case       =  "rgb"
      __snake_case       =  (16, 32, 96, 2_56)







      def       __lowerCAmelCase					(      self      :						Dict ,					__UpperCAmelCase      :						jax.random.KeyArray  )    ->FrozenDict:




       """simple docstring"""


       a		  =      (1, self.in_channels, self.sample_size, self.sample_size)
       a		  =      jnp.zeros(__UpperCAmelCase ,					dtype=jnp.floataa  )
       a		  =      jnp.ones((1,) ,					dtype=jnp.intaa  )
       a		  =      jnp.zeros((1, 1, self.cross_attention_dim) ,					dtype=jnp.floataa  )
       a		  =      (1, 3, self.sample_size * 8, self.sample_size * 8)
       a		  =      jnp.zeros(__UpperCAmelCase ,					dtype=jnp.floataa  )

       a  ,     a		  =      jax.random.split(__UpperCAmelCase  )
       a		  =      {'''params''': params_rng, '''dropout''': dropout_rng}

       return self.init(__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase  )["params"]







      def       __lowerCAmelCase					(      self      :						int  )    ->List[str]:




       """simple docstring"""


       a		  =      self.block_out_channels
       a		  =      block_out_channels[0] * 4

       # If `num_attention_heads` is not defined (which is the case for most models)
       # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
       # The reason for this behavior is to correct for incorrectly named variables that were introduced
       # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
       # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
       # which is why we correct for the naming here.
       a		  =      self.num_attention_heads or self.attention_head_dim

       # input
       a		  =      nn.Conv(
           block_out_channels[0] ,					kernel_size=(3, 3) ,					strides=(1, 1) ,					padding=((1, 1), (1, 1)) ,					dtype=self.dtype ,					)

       # time
       a		  =      FlaxTimesteps(
           block_out_channels[0] ,					flip_sin_to_cos=self.flip_sin_to_cos ,					freq_shift=self.config.freq_shift  )
       a		  =      FlaxTimestepEmbedding(__UpperCAmelCase ,					dtype=self.dtype  )

       a		  =      FlaxControlNetConditioningEmbedding(
           conditioning_embedding_channels=block_out_channels[0] ,					block_out_channels=self.conditioning_embedding_out_channels ,					)

       a		  =      self.only_cross_attention
       if isinstance(__UpperCAmelCase ,					__UpperCAmelCase  ):
        a		  =      (only_cross_attention,) * len(self.down_block_types  )

       if isinstance(__UpperCAmelCase ,					__UpperCAmelCase  ):
        a		  =      (num_attention_heads,) * len(self.down_block_types  )

       # down
       a		  =      []
       a		  =      []

       a		  =      block_out_channels[0]

       a		  =      nn.Conv(
           __UpperCAmelCase ,					kernel_size=(1, 1) ,					padding='''VALID''' ,					kernel_init=nn.initializers.zeros_init() ,					bias_init=nn.initializers.zeros_init() ,					dtype=self.dtype ,					)
       controlnet_down_blocks.append(__UpperCAmelCase  )

       for i, down_block_type in enumerate(self.down_block_types  ):
        a		  =      output_channel
        a		  =      block_out_channels[i]
        a		  =      i == len(__UpperCAmelCase  ) - 1

        if down_block_type == "CrossAttnDownBlock2D":
         a		  =      FlaxCrossAttnDownBlockaD(
             in_channels=__UpperCAmelCase ,					out_channels=__UpperCAmelCase ,					dropout=self.dropout ,					num_layers=self.layers_per_block ,					num_attention_heads=num_attention_heads[i] ,					add_downsample=not is_final_block ,					use_linear_projection=self.use_linear_projection ,					only_cross_attention=only_cross_attention[i] ,					dtype=self.dtype ,					)
        else:
         a		  =      FlaxDownBlockaD(
             in_channels=__UpperCAmelCase ,					out_channels=__UpperCAmelCase ,					dropout=self.dropout ,					num_layers=self.layers_per_block ,					add_downsample=not is_final_block ,					dtype=self.dtype ,					)

        down_blocks.append(__UpperCAmelCase  )

        for _ in range(self.layers_per_block  ):
         a		  =      nn.Conv(
             __UpperCAmelCase ,					kernel_size=(1, 1) ,					padding='''VALID''' ,					kernel_init=nn.initializers.zeros_init() ,					bias_init=nn.initializers.zeros_init() ,					dtype=self.dtype ,					)
         controlnet_down_blocks.append(__UpperCAmelCase  )

        if not is_final_block:
         a		  =      nn.Conv(
             __UpperCAmelCase ,					kernel_size=(1, 1) ,					padding='''VALID''' ,					kernel_init=nn.initializers.zeros_init() ,					bias_init=nn.initializers.zeros_init() ,					dtype=self.dtype ,					)
         controlnet_down_blocks.append(__UpperCAmelCase  )

       a		  =      down_blocks
       a		  =      controlnet_down_blocks

       # mid
       a		  =      block_out_channels[-1]
       a		  =      FlaxUNetMidBlockaDCrossAttn(
           in_channels=__UpperCAmelCase ,					dropout=self.dropout ,					num_attention_heads=num_attention_heads[-1] ,					use_linear_projection=self.use_linear_projection ,					dtype=self.dtype ,					)

       a		  =      nn.Conv(
           __UpperCAmelCase ,					kernel_size=(1, 1) ,					padding='''VALID''' ,					kernel_init=nn.initializers.zeros_init() ,					bias_init=nn.initializers.zeros_init() ,					dtype=self.dtype ,					)







      def __call__(      self      :						Any ,					__UpperCAmelCase      :						List[str] ,					__UpperCAmelCase      :						str ,					__UpperCAmelCase      :						int ,					__UpperCAmelCase      :						Tuple ,					__UpperCAmelCase      :						float = 1.0 ,					__UpperCAmelCase      :						bool = True ,					__UpperCAmelCase      :						bool = False ,					)    ->Union[FlaxControlNetOutput, Tuple]:




       """simple docstring"""


       a		  =      self.controlnet_conditioning_channel_order
       if channel_order == "bgr":
        a		  =      jnp.flip(__UpperCAmelCase ,					axis=1  )

       # 1. time
       if not isinstance(__UpperCAmelCase ,					jnp.ndarray  ):
        a		  =      jnp.array([timesteps] ,					dtype=jnp.intaa  )
       elif isinstance(__UpperCAmelCase ,					jnp.ndarray  ) and len(timesteps.shape  ) == 0:
        a		  =      timesteps.astype(dtype=jnp.floataa  )
        a		  =      jnp.expand_dims(__UpperCAmelCase ,					0  )

       a		  =      self.time_proj(__UpperCAmelCase  )
       a		  =      self.time_embedding(__UpperCAmelCase  )

       # 2. pre-process
       a		  =      jnp.transpose(__UpperCAmelCase ,					(0, 2, 3, 1)  )
       a		  =      self.conv_in(__UpperCAmelCase  )

       a		  =      jnp.transpose(__UpperCAmelCase ,					(0, 2, 3, 1)  )
       a		  =      self.controlnet_cond_embedding(__UpperCAmelCase  )
       sample += controlnet_cond

       # 3. down
       a		  =      (sample,)
       for down_block in self.down_blocks:
        if isinstance(__UpperCAmelCase ,					__UpperCAmelCase  ):
         a  ,     a		  =      down_block(__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					deterministic=not train  )
        else:
         a  ,     a		  =      down_block(__UpperCAmelCase ,					__UpperCAmelCase ,					deterministic=not train  )
        down_block_res_samples += res_samples

       # 4. mid
       a		  =      self.mid_block(__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					deterministic=not train  )

       # 5. contronet blocks
       a		  =      ()
       for down_block_res_sample, controlnet_block in zip(__UpperCAmelCase ,					self.controlnet_down_blocks  ):
        a		  =      controlnet_block(__UpperCAmelCase  )
        controlnet_down_block_res_samples += (down_block_res_sample,)

       a		  =      controlnet_down_block_res_samples

       a		  =      self.controlnet_mid_block(__UpperCAmelCase  )

       # 6. scaling
       a		  =      [sample * conditioning_scale for sample in down_block_res_samples]
       mid_block_res_sample *= conditioning_scale

       if not return_dict:
        return (down_block_res_samples, mid_block_res_sample)

       return FlaxControlNetOutput(
           down_block_res_samples=__UpperCAmelCase ,					mid_block_res_sample=__UpperCAmelCase  )







"""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        __SCREAMING_SNAKE_CASE	:


					'''simple docstring'''





					def __init__(       self							:  Optional[Any],  lowerCamelCase							:  Any,  lowerCamelCase							:  List[Any]=13,  lowerCamelCase							:  Any=10,  lowerCamelCase							:  Optional[Any]=3,  lowerCamelCase							:  Union[str, Any]=2,  lowerCamelCase							:  Dict=2,  lowerCamelCase							:  Tuple=2,  lowerCamelCase							:  List[str]=True,  lowerCamelCase							:  Optional[int]=True,  lowerCamelCase							:  Dict=32,  lowerCamelCase							:  Any=5,  lowerCamelCase							:  Dict=4,  lowerCamelCase							:  Any=37,  lowerCamelCase							:  Union[str, Any]="gelu",  lowerCamelCase							:  Dict=0.1,  lowerCamelCase							:  Union[str, Any]=0.1,  lowerCamelCase							:  Dict=10,  lowerCamelCase							:  str=0.02,  lowerCamelCase							:  List[Any]=0.9,  lowerCamelCase							:  List[Any]=None,  )->					str:
								lowerCamelCase__		:		List[str]			=parent
								lowerCamelCase__		:		Any			=batch_size
								lowerCamelCase__		:		str			=image_size
								lowerCamelCase__		:		Optional[Any]			=num_channels
								lowerCamelCase__		:		Optional[int]			=patch_size
								lowerCamelCase__		:		List[str]			=tubelet_size
								lowerCamelCase__		:		Optional[Any]			=num_frames
								lowerCamelCase__		:		Any			=is_training
								lowerCamelCase__		:		List[Any]			=use_labels
								lowerCamelCase__		:		Union[str, Any]			=hidden_size
								lowerCamelCase__		:		List[str]			=num_hidden_layers
								lowerCamelCase__		:		str			=num_attention_heads
								lowerCamelCase__		:		List[Any]			=intermediate_size
								lowerCamelCase__		:		Any			=hidden_act
								lowerCamelCase__		:		int			=hidden_dropout_prob
								lowerCamelCase__		:		Optional[int]			=attention_probs_dropout_prob
								lowerCamelCase__		:		Optional[Any]			=type_sequence_label_size
								lowerCamelCase__		:		int			=initializer_range
								lowerCamelCase__		:		Optional[Any]			=mask_ratio
								lowerCamelCase__		:		Any			=scope

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

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



					def 		snake_case       (       self							:  Dict     )->					Union[str, Any]:
								lowerCamelCase__		:		str			=floats_tensor(
								    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]     )

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

								lowerCamelCase__		:		Optional[Any]			=self.get_config()

								return config, pixel_values, labels



					def 		snake_case       (       self							:  Union[str, Any]     )->					Optional[int]:
								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=lowerCamelCase,  initializer_range=self.initializer_range,  )



					def 		snake_case       (       self							:  Dict,  lowerCamelCase							:  Tuple,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  Any     )->					Union[str, Any]:
								lowerCamelCase__		:		List[str]			=VideoMAEModel(config=lowerCamelCase     )
								model.to(lowerCamelCase     )
								model.eval()
								lowerCamelCase__		:		int			=model(lowerCamelCase     )
								self.parent.assertEqual(result.last_hidden_state.shape,  (self.batch_size, self.seq_length, self.hidden_size)     )



					def 		snake_case       (       self							:  Any,  lowerCamelCase							:  str,  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  str     )->					Dict:
								lowerCamelCase__		:		int			=VideoMAEForPreTraining(lowerCamelCase     )
								model.to(lowerCamelCase     )
								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
								lowerCamelCase__		:		Optional[int]			=torch.ones((self.num_masks,)     )
								lowerCamelCase__		:		List[str]			=torch.cat([mask, torch.zeros(self.seq_length - mask.size(0     )     )]     )
								lowerCamelCase__		:		int			=mask.expand(self.batch_size,  -1     ).bool()

								lowerCamelCase__		:		Any			=model(lowerCamelCase,  lowerCamelCase     )
								# model only returns predictions for masked patches
								lowerCamelCase__		:		Optional[int]			=mask.sum().item()
								lowerCamelCase__		:		Dict			=3 * self.tubelet_size * self.patch_size**2
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, num_masked_patches, decoder_num_labels)     )




					def 		snake_case       (       self							:  Optional[Any]     )->					Tuple:
								lowerCamelCase__		:		Tuple			=self.prepare_config_and_inputs()
								lowerCamelCase__	,   lowerCamelCase__	,   lowerCamelCase__		:		Dict			=config_and_inputs
								lowerCamelCase__		:		List[str]			={'''pixel_values''': pixel_values}
								return config, inputs_dict





@require_torch
class        __SCREAMING_SNAKE_CASE	(      lowerCAmelCase_		,   lowerCAmelCase_		,   unittest.TestCase   ):


					'''simple docstring'''





					_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 		snake_case       (       self							:  List[Any]     )->					Tuple:
								lowerCamelCase__		:		int			=VideoMAEModelTester(self     )
								lowerCamelCase__		:		Optional[int]			=ConfigTester(self,  config_class=lowerCamelCase,  has_text_modality=lowerCamelCase,  hidden_size=37     )



					def 		snake_case       (       self							:  Any,  lowerCamelCase							:  List[Any],  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  List[str]=False     )->					Tuple:
								lowerCamelCase__		:		str			=copy.deepcopy(lowerCamelCase     )

								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
											lowerCamelCase__		:		Any			=torch.ones((self.model_tester.num_masks,)     )
											lowerCamelCase__		:		Dict			=torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0     )     )]     )
											lowerCamelCase__		:		Optional[int]			=mask.expand(self.model_tester.batch_size,  -1     ).bool()
											lowerCamelCase__		:		int			=bool_masked_pos.to(lowerCamelCase     )

								if return_labels:
											if model_class in [
											    *get_values(lowerCamelCase     ),
											]:
														lowerCamelCase__		:		List[str]			=torch.zeros(
														    self.model_tester.batch_size,  dtype=torch.long,  device=lowerCamelCase     )

								return inputs_dict



					def 		snake_case       (       self							:  List[Any]     )->					int:
								self.config_tester.run_common_tests()



					@unittest.skip(reason='''VideoMAE does not use inputs_embeds'''     )
					def 		snake_case       (       self							:  List[str]     )->					Tuple:
								pass



					def 		snake_case       (       self							:  Union[str, Any]     )->					Union[str, Any]:
								lowerCamelCase__	,   lowerCamelCase__		:		Optional[int]			=self.model_tester.prepare_config_and_inputs_for_common()

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



					def 		snake_case       (       self							:  Optional[int]     )->					Optional[Any]:
								lowerCamelCase__	,   lowerCamelCase__		:		Any			=self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											lowerCamelCase__		:		Optional[int]			=model_class(lowerCamelCase     )
											lowerCamelCase__		:		Dict			=inspect.signature(model.forward     )
											# signature.parameters is an OrderedDict => so arg_names order is deterministic
											lowerCamelCase__		:		Tuple			=[*signature.parameters.keys()]

											lowerCamelCase__		:		List[str]			=['''pixel_values''']
											self.assertListEqual(arg_names[:1],  lowerCamelCase     )



					def 		snake_case       (       self							:  Tuple     )->					Optional[int]:
								lowerCamelCase__		:		Any			=self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCamelCase     )



					def 		snake_case       (       self							:  List[Any]     )->					Union[str, Any]:
								lowerCamelCase__		:		Optional[int]			=self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*lowerCamelCase     )



					@slow
					def 		snake_case       (       self							:  List[Any]     )->					Dict:
								for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											lowerCamelCase__		:		str			=VideoMAEModel.from_pretrained(lowerCamelCase     )
											self.assertIsNotNone(lowerCamelCase     )



					def 		snake_case       (       self							:  List[str]     )->					Optional[int]:
								if not self.has_attentions:
											pass

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

											for model_class in self.all_model_classes:
														lowerCamelCase__		:		Any			=self.model_tester.seq_length - self.model_tester.num_masks
														lowerCamelCase__		:		Any			=(
														    num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
														)

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

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

														self.assertListEqual(
														    list(attentions[0].shape[-3:]     ),  [self.model_tester.num_attention_heads, seq_len, seq_len],  )
														lowerCamelCase__		:		Union[str, Any]			=len(lowerCamelCase     )

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

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

														lowerCamelCase__		:		Optional[Any]			=outputs.attentions

														self.assertEqual(len(lowerCamelCase     ),  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 		snake_case       (       self							:  str     )->					int:
								def check_hidden_states_output(lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  List[str],  lowerCamelCase							:  Optional[Any]     ):
											lowerCamelCase__		:		List[Any]			=model_class(lowerCamelCase     )
											model.to(lowerCamelCase     )
											model.eval()

											with torch.no_grad():
														lowerCamelCase__		:		Optional[Any]			=model(**self._prepare_for_class(lowerCamelCase,  lowerCamelCase     )     )

											lowerCamelCase__		:		Dict			=outputs.hidden_states
											lowerCamelCase__		:		Any			=self.model_tester.num_hidden_layers + 1
											self.assertEqual(len(lowerCamelCase     ),  lowerCamelCase     )

											lowerCamelCase__		:		Any			=self.model_tester.seq_length - self.model_tester.num_masks
											lowerCamelCase__		:		str			=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],  )

								lowerCamelCase__	,   lowerCamelCase__		:		List[str]			=self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
											lowerCamelCase__		:		Union[str, Any]			=True
											check_hidden_states_output(lowerCamelCase,  lowerCamelCase,  lowerCamelCase     )

											# check that output_hidden_states also work using config
											del inputs_dict["output_hidden_states"]
											lowerCamelCase__		:		int			=True

											check_hidden_states_output(lowerCamelCase,  lowerCamelCase,  lowerCamelCase     )




					@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''     )
					def 		snake_case       (       self							:  Optional[int]     )->					int:
								pass






def      snake_case__   (		):
			"""simple docstring"""
			lowerCamelCase__		:		int			=hf_hub_download(
			    repo_id='''hf-internal-testing/spaghetti-video'''							,		filename='''eating_spaghetti.npy'''							,		repo_type='''dataset'''	)
			lowerCamelCase__		:		str			=np.load(__lowerCamelCase	)
			return list(__lowerCamelCase	)

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


					'''simple docstring'''





					@cached_property
					def 		snake_case       (       self							:  List[str]     )->					List[Any]:
								# logits were tested with a different mean and std, so we use the same here
								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 		snake_case       (       self							:  Optional[Any]     )->					Dict:
								lowerCamelCase__		:		str			=VideoMAEForVideoClassification.from_pretrained('''MCG-NJU/videomae-base-finetuned-kinetics'''     ).to(
								    lowerCamelCase     )

								lowerCamelCase__		:		Optional[Any]			=self.default_image_processor
								lowerCamelCase__		:		List[str]			=prepare_video()
								lowerCamelCase__		:		Union[str, Any]			=image_processor(lowerCamelCase,  return_tensors='''pt'''     ).to(lowerCamelCase     )

								# forward pass
								with torch.no_grad():
											lowerCamelCase__		:		Tuple			=model(**lowerCamelCase     )

								# verify the logits
								lowerCamelCase__		:		Union[str, Any]			=torch.Size((1, 400)     )
								self.assertEqual(outputs.logits.shape,  lowerCamelCase     )

								lowerCamelCase__		:		Tuple			=torch.tensor([0.3_669, -0.0_688, -0.2_421]     ).to(lowerCamelCase     )

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




					@slow
					def 		snake_case       (       self							:  Any     )->					Tuple:
								lowerCamelCase__		:		Tuple			=VideoMAEForPreTraining.from_pretrained('''MCG-NJU/videomae-base-short'''     ).to(lowerCamelCase     )

								lowerCamelCase__		:		Optional[int]			=self.default_image_processor
								lowerCamelCase__		:		Dict			=prepare_video()
								lowerCamelCase__		:		Dict			=image_processor(lowerCamelCase,  return_tensors='''pt'''     ).to(lowerCamelCase     )

								# add boolean mask, indicating which patches to mask
								lowerCamelCase__		:		str			=hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos''',  filename='''bool_masked_pos.pt'''     )
								lowerCamelCase__		:		Dict			=torch.load(lowerCamelCase     )

								# forward pass
								with torch.no_grad():
											lowerCamelCase__		:		Union[str, Any]			=model(**lowerCamelCase     )

								# verify the logits
								lowerCamelCase__		:		Dict			=torch.Size([1, 1408, 1536]     )
								lowerCamelCase__		:		Union[str, Any]			=torch.tensor(
								    [[0.7_994, 0.9_612, 0.8_508], [0.7_401, 0.8_958, 0.8_302], [0.5_862, 0.7_468, 0.7_325]],  device=lowerCamelCase     )
								self.assertEqual(outputs.logits.shape,  lowerCamelCase     )
								self.assertTrue(torch.allclose(outputs.logits[0, :3, :3],  lowerCamelCase,  atol=1E-4     )     )

								# verify the loss (`config.norm_pix_loss` = `True`)
								lowerCamelCase__		:		Optional[int]			=torch.tensor([0.5_142],  device=lowerCamelCase     )
								self.assertTrue(torch.allclose(outputs.loss,  lowerCamelCase,  atol=1E-4     )     )

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

								with torch.no_grad():
											lowerCamelCase__		:		Union[str, Any]			=model(**lowerCamelCase     )

								lowerCamelCase__		:		Union[str, Any]			=torch.tensor(torch.tensor([0.6_469]     ),  device=lowerCamelCase     )
								self.assertTrue(torch.allclose(outputs.loss,  lowerCamelCase,  atol=1E-4     )     )




