




from typing import Dict, Tuple, overload

import torch
import torch.types
from torch import nn

from . import residue_constants as rc
from .rigid_utils import Rigid, Rotation
from .tensor_utils import batched_gather




@overload
def     __SCREAMING_SNAKE_CASE      (     lowercase_    ,   lowercase_    ,   lowercase_					)				->		torch.Tensor:

    """simple docstring"""


    ...




@overload
def     __SCREAMING_SNAKE_CASE      (     lowercase_    ,   lowercase_    ,   lowercase_					)				->		Tuple[torch.Tensor, torch.Tensor]:

    """simple docstring"""


    ...




def     __SCREAMING_SNAKE_CASE      (     lowercase_    ,   lowercase_    ,   lowercase_					)				->		Union[str, Any]:

    """simple docstring"""


    __UpperCamelCase     						=  aatype == rc.restype_order['''G''']
    __UpperCamelCase     						=  rc.atom_order['''CA''']
    __UpperCamelCase     						=  rc.atom_order['''CB''']
    __UpperCamelCase     						=  torch.where(
        is_gly[..., None].expand(*((-1,) * len(is_gly.shape					))    ,   3					)    ,   all_atom_positions[..., ca_idx, :]    ,   all_atom_positions[..., cb_idx, :]    ,   )

    if all_atom_masks is not None:
        __UpperCamelCase     						=  torch.where(
            lowercase_    ,   all_atom_masks[..., ca_idx]    ,   all_atom_masks[..., cb_idx]    ,   )
        return pseudo_beta, pseudo_beta_mask
    else:
        return pseudo_beta




def     __SCREAMING_SNAKE_CASE      (     lowercase_    ,   lowercase_					)				->		torch.Tensor:

    """simple docstring"""


    __UpperCamelCase     						=  batched_gather(
        lowercase_    ,   batch['''residx_atom37_to_atom14''']    ,   dim=-2    ,   no_batch_dims=len(atomaa.shape[:-2]					)    ,   )

    __UpperCamelCase     						=  atomaa_data * batch['''atom37_atom_exists'''][..., None]

    return atomaa_data




def     __SCREAMING_SNAKE_CASE      (     lowercase_					)				->		torch.Tensor:

    """simple docstring"""


    __UpperCamelCase     						=  template_feats['''template_aatype''']
    __UpperCamelCase     						=  template_feats['''template_torsion_angles_sin_cos''']
    __UpperCamelCase     						=  template_feats['''template_alt_torsion_angles_sin_cos''']
    __UpperCamelCase     						=  template_feats['''template_torsion_angles_mask''']
    __UpperCamelCase     						=  torch.cat(
        [
            nn.functional.one_hot(lowercase_    ,   22					),
            torsion_angles_sin_cos.reshape(*torsion_angles_sin_cos.shape[:-2]    ,   14					),
            alt_torsion_angles_sin_cos.reshape(*alt_torsion_angles_sin_cos.shape[:-2]    ,   14					),
            torsion_angles_mask,
        ]    ,   dim=-1    ,   )

    return template_angle_feat




def     __SCREAMING_SNAKE_CASE      (     lowercase_    ,   lowercase_    ,   lowercase_    ,   lowercase_    ,   lowercase_ = False    ,   lowercase_ = 1E-20    ,   lowercase_ = 1E8    ,   )				->		torch.Tensor:

    """simple docstring"""


    __UpperCamelCase     						=  batch['''template_pseudo_beta_mask''']
    __UpperCamelCase     						=  template_mask[..., None] * template_mask[..., None, :]

    # Compute distogram (this seems to differ slightly from Alg. 5)
    __UpperCamelCase     						=  batch['''template_pseudo_beta''']
    __UpperCamelCase     						=  torch.sum((tpb[..., None, :] - tpb[..., None, :, :]) ** 2    ,   dim=-1    ,   keepdim=lowercase_					)
    __UpperCamelCase     						=  torch.linspace(lowercase_    ,   lowercase_    ,   lowercase_    ,   device=tpb.device					) ** 2
    __UpperCamelCase     						=  torch.cat([lower[1:], lower.new_tensor([inf]					)]    ,   dim=-1					)
    __UpperCamelCase     						=  ((dgram > lower) * (dgram < upper)).type(dgram.dtype					)

    __UpperCamelCase     						=  [dgram, template_mask_ad[..., None]]

    __UpperCamelCase     						=  nn.functional.one_hot(
        batch['''template_aatype''']    ,   rc.restype_num + 2    ,   )

    __UpperCamelCase     						=  batch['''template_aatype'''].shape[-1]
    to_concat.append(aatype_one_hot[..., None, :, :].expand(*aatype_one_hot.shape[:-2]    ,   lowercase_    ,   -1    ,   -1					)					)
    to_concat.append(aatype_one_hot[..., None, :].expand(*aatype_one_hot.shape[:-2]    ,   -1    ,   lowercase_    ,   -1					)					)

    __UpperCamelCase						,					__UpperCamelCase						,					__UpperCamelCase     						=  [rc.atom_order[a] for a in ['''N''', '''CA''', '''C''']]
    __UpperCamelCase     						=  Rigid.make_transform_from_reference(
        n_xyz=batch['''template_all_atom_positions'''][..., n, :]    ,   ca_xyz=batch['''template_all_atom_positions'''][..., ca, :]    ,   c_xyz=batch['''template_all_atom_positions'''][..., c, :]    ,   eps=lowercase_    ,   )
    __UpperCamelCase     						=  rigids.get_trans()[..., None, :, :]
    __UpperCamelCase     						=  rigids[..., None].invert_apply(lowercase_					)

    __UpperCamelCase     						=  torch.rsqrt(eps + torch.sum(rigid_vec**2    ,   dim=-1					)					)

    __UpperCamelCase     						=  batch['''template_all_atom_mask''']
    __UpperCamelCase     						=  t_aa_masks[..., n] * t_aa_masks[..., ca] * t_aa_masks[..., c]
    __UpperCamelCase     						=  template_mask[..., None] * template_mask[..., None, :]

    __UpperCamelCase     						=  inv_distance_scalar * template_mask_ad
    __UpperCamelCase     						=  rigid_vec * inv_distance_scalar[..., None]

    if not use_unit_vector:
        __UpperCamelCase     						=  unit_vector * 0.0

    to_concat.extend(torch.unbind(unit_vector[..., None, :]    ,   dim=-1					)					)
    to_concat.append(template_mask_ad[..., None]					)

    __UpperCamelCase     						=  torch.cat(lowercase_    ,   dim=-1					)
    __UpperCamelCase     						=  act * template_mask_ad[..., None]

    return act




def     __SCREAMING_SNAKE_CASE      (     lowercase_					)				->		torch.Tensor:

    """simple docstring"""


    __UpperCamelCase     						=  nn.functional.one_hot(batch['''extra_msa''']    ,   23					)
    __UpperCamelCase     						=  [
        msa_ahot,
        batch['''extra_has_deletion'''].unsqueeze(-1					),
        batch['''extra_deletion_value'''].unsqueeze(-1					),
    ]
    return torch.cat(lowercase_    ,   dim=-1					)




def     __SCREAMING_SNAKE_CASE      (     lowercase_    ,   lowercase_    ,   lowercase_    ,   lowercase_    ,   )				->		Rigid:

    """simple docstring"""


    __UpperCamelCase     						=  rrgdf[aatype, ...]

    # [*, N, 8] transformations, i.e.
    #   One [*, N, 8, 3, 3] rotation matrix and
    #   One [*, N, 8, 3]    translation matrix
    __UpperCamelCase     						=  r.from_tensor_axa(lowercase_					)

    __UpperCamelCase     						=  alpha.new_zeros((*((1,) * len(alpha.shape[:-1]					)), 2)					)
    __UpperCamelCase     						=  1

    # [*, N, 8, 2]
    __UpperCamelCase     						=  torch.cat([bb_rot.expand(*alpha.shape[:-2]    ,   -1    ,   -1					), alpha]    ,   dim=-2					)

    # [*, N, 8, 3, 3]
    # Produces rotation matrices of the form:
    # [
    #   [1, 0  , 0  ],
    #   [0, a_2,-a_1],
    #   [0, a_1, a_2]
    # ]
    # This follows the original code rather than the supplement, which uses
    # different indices.

    __UpperCamelCase     						=  alpha.new_zeros(default_r.get_rots().get_rot_mats().shape					)
    __UpperCamelCase     						=  1
    __UpperCamelCase     						=  alpha[..., 1]
    __UpperCamelCase     						=  -alpha[..., 0]
    __UpperCamelCase     						=  alpha

    __UpperCamelCase     						=  default_r.compose(Rigid(Rotation(rot_mats=lowercase_					)    ,   lowercase_					)					)

    __UpperCamelCase     						=  all_frames[..., 5]
    __UpperCamelCase     						=  all_frames[..., 6]
    __UpperCamelCase     						=  all_frames[..., 7]

    __UpperCamelCase     						=  all_frames[..., 4]
    __UpperCamelCase     						=  chia_frame_to_bb.compose(lowercase_					)
    __UpperCamelCase     						=  chia_frame_to_bb.compose(lowercase_					)
    __UpperCamelCase     						=  chia_frame_to_bb.compose(lowercase_					)

    __UpperCamelCase     						=  Rigid.cat(
        [
            all_frames[..., :5],
            chia_frame_to_bb.unsqueeze(-1					),
            chia_frame_to_bb.unsqueeze(-1					),
            chia_frame_to_bb.unsqueeze(-1					),
        ]    ,   dim=-1    ,   )

    __UpperCamelCase     						=  r[..., None].compose(lowercase_					)

    return all_frames_to_global




def     __SCREAMING_SNAKE_CASE      (     lowercase_    ,   lowercase_    ,   lowercase_    ,   lowercase_    ,   lowercase_    ,   lowercase_    ,   )				->		torch.Tensor:

    """simple docstring"""


    __UpperCamelCase     						=  group_idx[aatype, ...]

    # [*, N, 14, 8]
    __UpperCamelCase     						=  nn.functional.one_hot(
        lowercase_    ,   num_classes=default_frames.shape[-3]    ,   )

    # [*, N, 14, 8]
    __UpperCamelCase     						=  r[..., None, :] * group_mask_one_hot

    # [*, N, 14]
    __UpperCamelCase     						=  t_atoms_to_global.map_tensor_fn(lambda lowercase_					: torch.sum(lowercase_    ,   dim=-1					)					)

    # [*, N, 14, 1]
    __UpperCamelCase     						=  atom_mask[aatype, ...].unsqueeze(-1					)

    # [*, N, 14, 3]
    __UpperCamelCase     						=  lit_positions[aatype, ...]
    __UpperCamelCase     						=  t_atoms_to_global.apply(lowercase_					)
    __UpperCamelCase     						=  pred_positions * atom_mask

    return pred_positions










import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VivitConfig
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


if is_torch_available():
							import torch
							from torch import nn

							from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VivitForVideoClassification, VivitModel
							from transformers.models.vivit.modeling_vivit import VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from transformers import VivitImageProcessor

class   A__							:




						def __init__(				self    :					str ,			_a    :					List[str] ,			_a    :					List[str]=2 ,			_a    :					Dict=True ,			_a    :					Tuple=True ,			_a    :					Any=10 ,			_a    :					str=10 ,			_a    :					int=8 ,			_a    :					Tuple=[2, 4, 4] ,			_a    :					Optional[int]=3 ,			_a    :					List[str]=768 ,			_a    :					int=5 ,			_a    :					Any=4 ,			_a    :					Tuple=37 ,			_a    :					Tuple="gelu_fast" ,			_a    :					Dict=0.0 ,			_a    :					Optional[Any]=0.0 ,			_a    :					Optional[Any]=0.02 ,			_a    :					str=1E-06 ,			_a    :					Tuple=True ,			_a    :					List[str]=None ,			)							->   Dict:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =parent
										_SCREAMING_SNAKE_CASE	     =batch_size
										_SCREAMING_SNAKE_CASE	     =is_training
										_SCREAMING_SNAKE_CASE	     =use_labels
										_SCREAMING_SNAKE_CASE	     =num_labels
										_SCREAMING_SNAKE_CASE	     =image_size
										_SCREAMING_SNAKE_CASE	     =num_frames
										_SCREAMING_SNAKE_CASE	     =tubelet_size
										_SCREAMING_SNAKE_CASE	     =num_channels
										_SCREAMING_SNAKE_CASE	     =hidden_size
										_SCREAMING_SNAKE_CASE	     =num_hidden_layers
										_SCREAMING_SNAKE_CASE	     =num_attention_heads
										_SCREAMING_SNAKE_CASE	     =intermediate_size
										_SCREAMING_SNAKE_CASE	     =hidden_act
										_SCREAMING_SNAKE_CASE	     =hidden_dropout_prob
										_SCREAMING_SNAKE_CASE	     =attention_probs_dropout_prob
										_SCREAMING_SNAKE_CASE	     =initializer_range
										_SCREAMING_SNAKE_CASE	     =layer_norm_eps
										_SCREAMING_SNAKE_CASE	     =qkv_bias
										_SCREAMING_SNAKE_CASE	     =scope

										_SCREAMING_SNAKE_CASE	     =(
										    (self.image_size // self.tubelet_size[2])
										    * (self.image_size // self.tubelet_size[1])
										    * (self.num_frames // self.tubelet_size[0])
										) + 1  # CLS token




						def 							__UpperCamelCase (				self    :					List[str]     )							->   str:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =floats_tensor(
										    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]     )

										_SCREAMING_SNAKE_CASE	     =None
										if self.use_labels:
														_SCREAMING_SNAKE_CASE	     =ids_tensor([self.batch_size] ,			self.num_labels     )

										_SCREAMING_SNAKE_CASE	     =self.get_config()

										return config, pixel_values, labels




						def 							__UpperCamelCase (				self    :					Optional[int]     )							->   List[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =VivitConfig(
										    num_frames=self.num_frames ,			image_size=self.image_size ,			tubelet_size=self.tubelet_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 ,			initializer_range=self.initializer_range ,			layer_norm_eps=self.layer_norm_eps ,			qkv_bias=self.qkv_bias ,			)
										_SCREAMING_SNAKE_CASE	     =self.num_labels
										return config




						def 							__UpperCamelCase (				self    :					Union[str, Any] ,			_a    :					Optional[int] ,			_a    :					Dict ,			_a    :					Tuple     )							->   Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =VivitModel(config=_a     )
										model.to(_a     )
										model.eval()
										_SCREAMING_SNAKE_CASE	     =model(_a     )
										self.parent.assertEqual(result.last_hidden_state.shape ,			(self.batch_size, self.seq_length, self.hidden_size)     )




						def 							__UpperCamelCase (				self    :					Union[str, Any] ,			_a    :					Optional[Any] ,			_a    :					Any ,			_a    :					Any     )							->   int:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =VivitForVideoClassification(_a     )
										model.to(_a     )
										model.eval()

										_SCREAMING_SNAKE_CASE	     =model(_a     )

										# verify the logits shape
										_SCREAMING_SNAKE_CASE	     =torch.Size((self.batch_size, self.num_labels)     )
										self.parent.assertEqual(result.logits.shape ,			_a     )




						def 							__UpperCamelCase (				self    :					Tuple     )							->   List[str]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =self.prepare_config_and_inputs()
										_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =config_and_inputs
										_SCREAMING_SNAKE_CASE	     ={'''pixel_values''': pixel_values}
										return config, inputs_dict

@require_torch
class   A__							(						UpperCamelCase__      ,   unittest.TestCase     ):
						UpperCAmelCase									=  (VivitModel, VivitForVideoClassification) if is_torch_available() else ()

						UpperCAmelCase									=  False
						UpperCAmelCase									=  False
						UpperCAmelCase									=  False
						UpperCAmelCase									=  False




						def 							__UpperCamelCase (				self    :					Any     )							->   Any:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =VivitModelTester(self     )
										_SCREAMING_SNAKE_CASE	     =ConfigTester(self ,			config_class=_a ,			has_text_modality=_a ,			hidden_size=37     )




						def 							__UpperCamelCase (				self    :					List[str] ,			_a    :					List[str] ,			_a    :					Union[str, Any] ,			_a    :					Dict=False     )							->   List[str]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =copy.deepcopy(_a     )

										if return_labels:
														if model_class in get_values(_a     ):
																		_SCREAMING_SNAKE_CASE	     =torch.zeros(
																		    self.model_tester.batch_size ,			dtype=torch.long ,			device=_a     )

										return inputs_dict




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


										"""simple docstring"""





										self.config_tester.run_common_tests()




						@unittest.skip(reason='''Vivit does not use inputs_embeds'''     )
						def 							__UpperCamelCase (				self    :					Optional[Any]     )							->   Any:


										"""simple docstring"""





										pass




						def 							__UpperCamelCase (				self    :					str     )							->   Dict:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =self.model_tester.prepare_config_and_inputs_for_common()

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




						def 							__UpperCamelCase (				self    :					int     )							->   str:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =self.model_tester.prepare_config_and_inputs_for_common()

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

														_SCREAMING_SNAKE_CASE	     =['''pixel_values''', '''head_mask''']
														self.assertListEqual(arg_names[:2] ,			_a     )




						def 							__UpperCamelCase (				self    :					Tuple     )							->   List[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_a     )




						def 							__UpperCamelCase (				self    :					str     )							->   Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_video_classification(*_a     )




						@slow
						def 							__UpperCamelCase (				self    :					str     )							->   Optional[Any]:


										"""simple docstring"""





										for model_name in VIVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_SCREAMING_SNAKE_CASE	     =VivitModel.from_pretrained(_a     )
														self.assertIsNotNone(_a     )




						def 							__UpperCamelCase (				self    :					List[Any]     )							->   Tuple:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =self.model_tester.prepare_config_and_inputs_for_common()
										_SCREAMING_SNAKE_CASE	     =True

										for model_class in self.all_model_classes:
														_SCREAMING_SNAKE_CASE	     =self.model_tester.seq_length

														_SCREAMING_SNAKE_CASE	     =True
														_SCREAMING_SNAKE_CASE	     =False
														_SCREAMING_SNAKE_CASE	     =True
														_SCREAMING_SNAKE_CASE	     =model_class(_a     )
														model.to(_a     )
														model.eval()
														with torch.no_grad():
																		_SCREAMING_SNAKE_CASE	     =model(**self._prepare_for_class(_a ,			_a     )     )
														_SCREAMING_SNAKE_CASE	     =outputs.attentions
														self.assertEqual(len(_a     ) ,			self.model_tester.num_hidden_layers     )

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

														self.assertListEqual(
														    list(attentions[0].shape[-3:]     ) ,			[self.model_tester.num_attention_heads, seq_len, seq_len] ,			)
														_SCREAMING_SNAKE_CASE	     =len(_a     )

														# Check attention is always last and order is fine
														_SCREAMING_SNAKE_CASE	     =True
														_SCREAMING_SNAKE_CASE	     =True
														_SCREAMING_SNAKE_CASE	     =model_class(_a     )
														model.to(_a     )
														model.eval()
														with torch.no_grad():
																		_SCREAMING_SNAKE_CASE	     =model(**self._prepare_for_class(_a ,			_a     )     )

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

														_SCREAMING_SNAKE_CASE	     =outputs.attentions

														self.assertEqual(len(_a     ) ,			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 							__UpperCamelCase (				self    :					Any     )							->   Optional[int]:


										"""simple docstring"""





										def check_hidden_states_output(_a    :					Dict ,			_a    :					Optional[Any] ,			_a    :					List[Any]     ):
														_SCREAMING_SNAKE_CASE	     =model_class(_a     )
														model.to(_a     )
														model.eval()

														with torch.no_grad():
																		_SCREAMING_SNAKE_CASE	     =model(**self._prepare_for_class(_a ,			_a     )     )

														_SCREAMING_SNAKE_CASE	     =outputs.hidden_states
														_SCREAMING_SNAKE_CASE	     =self.model_tester.num_hidden_layers + 1
														self.assertEqual(len(_a     ) ,			_a     )

														_SCREAMING_SNAKE_CASE	     =self.model_tester.seq_length

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

										_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														_SCREAMING_SNAKE_CASE	     =True
														check_hidden_states_output(_a ,			_a ,			_a     )

														# check that output_hidden_states also work using config
														del inputs_dict["output_hidden_states"]
														_SCREAMING_SNAKE_CASE	     =True

														check_hidden_states_output(_a ,			_a ,			_a     )




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


										"""simple docstring"""





										pass







def   lowerCamelCase(      ):
				_SCREAMING_SNAKE_CASE	     =hf_hub_download(
				    repo_id='''hf-internal-testing/spaghetti-video'''		,filename='''eating_spaghetti_32_frames.npy'''		,repo_type='''dataset''')
				_SCREAMING_SNAKE_CASE	     =np.load(a__)
				return list(a__)


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




						@cached_property
						def 							__UpperCamelCase (				self    :					Union[str, Any]     )							->   List[Any]:


										"""simple docstring"""





										return VivitImageProcessor() if is_vision_available() else None




						@slow
						def 							__UpperCamelCase (				self    :					Optional[Any]     )							->   str:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =VivitForVideoClassification.from_pretrained('''google/vivit-b-16x2-kinetics400'''     ).to(_a     )

										_SCREAMING_SNAKE_CASE	     =self.default_image_processor
										_SCREAMING_SNAKE_CASE	     =prepare_video()
										_SCREAMING_SNAKE_CASE	     =image_processor(_a ,			return_tensors='''pt'''     ).to(_a     )

										# forward pass
										with torch.no_grad():
														_SCREAMING_SNAKE_CASE	     =model(**_a     )

										# verify the logits
										_SCREAMING_SNAKE_CASE	     =torch.Size((1, 400)     )
										self.assertEqual(outputs.logits.shape ,			_a     )

										# taken from original model
										_SCREAMING_SNAKE_CASE	     =torch.tensor([-1.05_43, 2.07_64, -0.21_04, 0.44_39, -0.96_58]     ).to(_a     )

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