


'''simple docstring'''
import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp







def 			a_							(      UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_ = 4_0_9_6   ):
							A_						,				A_						,				A_ 						=		key.shape[-3:]
							A_ 						=		value.shape[-1]
							A_ 						=		min(UpperCamelCase_	,     UpperCamelCase_   )
							A_ 						=		query / jnp.sqrt(UpperCamelCase_   )

							@functools.partial(jax.checkpoint	,     prevent_cse=UpperCamelCase_   )
							def summarize_chunk(UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_   ):
														A_ 						=		jnp.einsum("...qhd,...khd->...qhk"	,     UpperCamelCase_	,     UpperCamelCase_	,     precision=UpperCamelCase_   )

														A_ 						=		jnp.max(UpperCamelCase_	,     axis=-1	,     keepdims=UpperCamelCase_   )
														A_ 						=		jax.lax.stop_gradient(UpperCamelCase_   )
														A_ 						=		jnp.exp(attn_weights - max_score   )

														A_ 						=		jnp.einsum("...vhf,...qhv->...qhf"	,     UpperCamelCase_	,     UpperCamelCase_	,     precision=UpperCamelCase_   )
														A_ 						=		jnp.einsum("...qhk->...qh"	,     UpperCamelCase_   )

														return (exp_values, exp_weights.sum(axis=-1   ), max_score)

							def chunk_scanner(UpperCamelCase_   ):
														# julienne key array
														A_ 						=		jax.lax.dynamic_slice(
														    operand=UpperCamelCase_	,     start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]	,     slice_sizes=list(key.shape[:-3]   ) + [key_chunk_size, num_heads, k_features]	,     )

														# julienne value array
														A_ 						=		jax.lax.dynamic_slice(
														    operand=UpperCamelCase_	,     start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]	,     slice_sizes=list(value.shape[:-3]   ) + [key_chunk_size, num_heads, v_features]	,     )

														return summarize_chunk(UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_   )

							A_						,				A_						,				A_ 						=		jax.lax.map(f=UpperCamelCase_	,     xs=jnp.arange(0	,     UpperCamelCase_	,     UpperCamelCase_   )   )

							A_ 						=		jnp.max(UpperCamelCase_	,     axis=0	,     keepdims=UpperCamelCase_   )
							A_ 						=		jnp.exp(chunk_max - global_max   )

							chunk_values *= jnp.expand_dims(UpperCamelCase_	,     axis=-1   )
							chunk_weights *= max_diffs

							A_ 						=		chunk_values.sum(axis=0   )
							A_ 						=		jnp.expand_dims(UpperCamelCase_	,     -1   ).sum(axis=0   )

							return all_values / all_weights







def 			a_							(      UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_	,     UpperCamelCase_=jax.lax.Precision.HIGHEST	,     UpperCamelCase_ = 1_0_2_4	,     UpperCamelCase_ = 4_0_9_6   ):
							A_						,				A_						,				A_ 						=		query.shape[-3:]

							def chunk_scanner(UpperCamelCase_	,     UpperCamelCase_   ):
														# julienne query array
														A_ 						=		jax.lax.dynamic_slice(
														    operand=UpperCamelCase_	,     start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]	,     slice_sizes=list(query.shape[:-3]   ) + [min(UpperCamelCase_	,     UpperCamelCase_   ), num_heads, q_features]	,     )

														return (
														    chunk_idx + query_chunk_size,  # unused ignore it
														    _query_chunk_attention(
														        query=UpperCamelCase_	,     key=UpperCamelCase_	,     value=UpperCamelCase_	,     precision=UpperCamelCase_	,     key_chunk_size=UpperCamelCase_   ),
														)

							A_						,				A_ 						=		jax.lax.scan(
							    f=UpperCamelCase_	,     init=0	,     xs=UpperCamelCase_	,     length=math.ceil(num_q / query_chunk_size   )   )

							return jnp.concatenate(UpperCamelCase_	,     axis=-3   )  # fuse the chunked result back






class 				__lowerCAmelCase      (      nn.Module	):

						"""simple docstring"""






						_UpperCAmelCase		:					int
						_UpperCAmelCase		:					int      	=8
						_UpperCAmelCase		:					int      	=64
						_UpperCAmelCase		:					float      	=0.0
						_UpperCAmelCase		:					bool      	=False
						_UpperCAmelCase		:					jnp.dtype      	=jnp.floataa






						def 							_UpperCAmelCase					(     self       :			Dict  ):
													A_ 						=		self.dim_head * self.heads
													A_ 						=		self.dim_head**-0.5

													# Weights were exported with old names {to_q, to_k, to_v, to_out}
													A_ 						=		nn.Dense(lowerCAmelCase   ,     use_bias=lowerCAmelCase   ,     dtype=self.dtype   ,     name="to_q"  )
													A_ 						=		nn.Dense(lowerCAmelCase   ,     use_bias=lowerCAmelCase   ,     dtype=self.dtype   ,     name="to_k"  )
													A_ 						=		nn.Dense(lowerCAmelCase   ,     use_bias=lowerCAmelCase   ,     dtype=self.dtype   ,     name="to_v"  )

													A_ 						=		nn.Dense(self.query_dim   ,     dtype=self.dtype   ,     name="to_out_0"  )
													A_ 						=		nn.Dropout(rate=self.dropout  )






						def 							_UpperCAmelCase					(     self       :			Any   ,     lowerCAmelCase       :			Any  ):
													A_						,				A_						,				A_ 						=		tensor.shape
													A_ 						=		self.heads
													A_ 						=		tensor.reshape(lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   ,     dim // head_size  )
													A_ 						=		jnp.transpose(lowerCAmelCase   ,     (0, 2, 1, 3)  )
													A_ 						=		tensor.reshape(batch_size * head_size   ,     lowerCAmelCase   ,     dim // head_size  )
													return tensor






						def 							_UpperCAmelCase					(     self       :			Any   ,     lowerCAmelCase       :			Optional[int]  ):
													A_						,				A_						,				A_ 						=		tensor.shape
													A_ 						=		self.heads
													A_ 						=		tensor.reshape(batch_size // head_size   ,     lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase  )
													A_ 						=		jnp.transpose(lowerCAmelCase   ,     (0, 2, 1, 3)  )
													A_ 						=		tensor.reshape(batch_size // head_size   ,     lowerCAmelCase   ,     dim * head_size  )
													return tensor






						def __call__(     self       :			List[str]   ,     lowerCAmelCase       :			Dict   ,     lowerCAmelCase       :			Tuple=None   ,     lowerCAmelCase       :			Optional[Any]=True  ):
													A_ 						=		hidden_states if context is None else context

													A_ 						=		self.query(lowerCAmelCase  )
													A_ 						=		self.key(lowerCAmelCase  )
													A_ 						=		self.value(lowerCAmelCase  )

													A_ 						=		self.reshape_heads_to_batch_dim(lowerCAmelCase  )
													A_ 						=		self.reshape_heads_to_batch_dim(lowerCAmelCase  )
													A_ 						=		self.reshape_heads_to_batch_dim(lowerCAmelCase  )

													if self.use_memory_efficient_attention:
																				A_ 						=		query_states.transpose(1   ,     0   ,     2  )
																				A_ 						=		key_states.transpose(1   ,     0   ,     2  )
																				A_ 						=		value_states.transpose(1   ,     0   ,     2  )

																				# this if statement create a chunk size for each layer of the unet
																				# the chunk size is equal to the query_length dimension of the deepest layer of the unet

																				A_ 						=		query_states.shape[-3]
																				if flatten_latent_dim % 64 == 0:
																											A_ 						=		int(flatten_latent_dim / 64  )
																				elif flatten_latent_dim % 16 == 0:
																											A_ 						=		int(flatten_latent_dim / 16  )
																				elif flatten_latent_dim % 4 == 0:
																											A_ 						=		int(flatten_latent_dim / 4  )
																				else:
																											A_ 						=		int(lowerCAmelCase  )

																				A_ 						=		jax_memory_efficient_attention(
																				    lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   ,     query_chunk_size=lowerCAmelCase   ,     key_chunk_size=40_96 * 4  )

																				A_ 						=		hidden_states.transpose(1   ,     0   ,     2  )
													else:
																				# compute attentions
																				A_ 						=		jnp.einsum("b i d, b j d->b i j"   ,     lowerCAmelCase   ,     lowerCAmelCase  )
																				A_ 						=		attention_scores * self.scale
																				A_ 						=		nn.softmax(lowerCAmelCase   ,     axis=2  )

																				# attend to values
																				A_ 						=		jnp.einsum("b i j, b j d -> b i d"   ,     lowerCAmelCase   ,     lowerCAmelCase  )

													A_ 						=		self.reshape_batch_dim_to_heads(lowerCAmelCase  )
													A_ 						=		self.proj_attn(lowerCAmelCase  )
													return self.dropout_layer(lowerCAmelCase   ,     deterministic=lowerCAmelCase  )






class 				__lowerCAmelCase      (      nn.Module	):

						"""simple docstring"""






						_UpperCAmelCase		:					int
						_UpperCAmelCase		:					int
						_UpperCAmelCase		:					int
						_UpperCAmelCase		:					float      	=0.0
						_UpperCAmelCase		:					bool      	=False
						_UpperCAmelCase		:					jnp.dtype      	=jnp.floataa
						_UpperCAmelCase		:					bool      	=False






						def 							_UpperCAmelCase					(     self       :			List[Any]  ):
													# self attention (or cross_attention if only_cross_attention is True)
													A_ 						=		FlaxAttention(
													    self.dim   ,     self.n_heads   ,     self.d_head   ,     self.dropout   ,     self.use_memory_efficient_attention   ,     dtype=self.dtype  )
													# cross attention
													A_ 						=		FlaxAttention(
													    self.dim   ,     self.n_heads   ,     self.d_head   ,     self.dropout   ,     self.use_memory_efficient_attention   ,     dtype=self.dtype  )
													A_ 						=		FlaxFeedForward(dim=self.dim   ,     dropout=self.dropout   ,     dtype=self.dtype  )
													A_ 						=		nn.LayerNorm(epsilon=1e-5   ,     dtype=self.dtype  )
													A_ 						=		nn.LayerNorm(epsilon=1e-5   ,     dtype=self.dtype  )
													A_ 						=		nn.LayerNorm(epsilon=1e-5   ,     dtype=self.dtype  )
													A_ 						=		nn.Dropout(rate=self.dropout  )






						def __call__(     self       :			List[Any]   ,     lowerCAmelCase       :			Dict   ,     lowerCAmelCase       :			Tuple   ,     lowerCAmelCase       :			Tuple=True  ):
													# self attention
													A_ 						=		hidden_states
													if self.only_cross_attention:
																				A_ 						=		self.attna(self.norma(lowerCAmelCase  )   ,     lowerCAmelCase   ,     deterministic=lowerCAmelCase  )
													else:
																				A_ 						=		self.attna(self.norma(lowerCAmelCase  )   ,     deterministic=lowerCAmelCase  )
													A_ 						=		hidden_states + residual

													# cross attention
													A_ 						=		hidden_states
													A_ 						=		self.attna(self.norma(lowerCAmelCase  )   ,     lowerCAmelCase   ,     deterministic=lowerCAmelCase  )
													A_ 						=		hidden_states + residual

													# feed forward
													A_ 						=		hidden_states
													A_ 						=		self.ff(self.norma(lowerCAmelCase  )   ,     deterministic=lowerCAmelCase  )
													A_ 						=		hidden_states + residual

													return self.dropout_layer(lowerCAmelCase   ,     deterministic=lowerCAmelCase  )






class 				__lowerCAmelCase      (      nn.Module	):

						"""simple docstring"""






						_UpperCAmelCase		:					int
						_UpperCAmelCase		:					int
						_UpperCAmelCase		:					int
						_UpperCAmelCase		:					int      	=1
						_UpperCAmelCase		:					float      	=0.0
						_UpperCAmelCase		:					bool      	=False
						_UpperCAmelCase		:					bool      	=False
						_UpperCAmelCase		:					jnp.dtype      	=jnp.floataa
						_UpperCAmelCase		:					bool      	=False






						def 							_UpperCAmelCase					(     self       :			Optional[int]  ):
													A_ 						=		nn.GroupNorm(num_groups=32   ,     epsilon=1e-5  )

													A_ 						=		self.n_heads * self.d_head
													if self.use_linear_projection:
																				A_ 						=		nn.Dense(lowerCAmelCase   ,     dtype=self.dtype  )
													else:
																				A_ 						=		nn.Conv(
																				    lowerCAmelCase   ,     kernel_size=(1, 1)   ,     strides=(1, 1)   ,     padding="VALID"   ,     dtype=self.dtype   ,     )

													A_ 						=		[
													    FlaxBasicTransformerBlock(
													        lowerCAmelCase   ,     self.n_heads   ,     self.d_head   ,     dropout=self.dropout   ,     only_cross_attention=self.only_cross_attention   ,     dtype=self.dtype   ,     use_memory_efficient_attention=self.use_memory_efficient_attention   ,     )
													    for _ in range(self.depth  )
													]

													if self.use_linear_projection:
																				A_ 						=		nn.Dense(lowerCAmelCase   ,     dtype=self.dtype  )
													else:
																				A_ 						=		nn.Conv(
																				    lowerCAmelCase   ,     kernel_size=(1, 1)   ,     strides=(1, 1)   ,     padding="VALID"   ,     dtype=self.dtype   ,     )

													A_ 						=		nn.Dropout(rate=self.dropout  )






						def __call__(     self       :			Union[str, Any]   ,     lowerCAmelCase       :			Union[str, Any]   ,     lowerCAmelCase       :			int   ,     lowerCAmelCase       :			List[str]=True  ):
													A_						,				A_						,				A_						,				A_ 						=		hidden_states.shape
													A_ 						=		hidden_states
													A_ 						=		self.norm(lowerCAmelCase  )
													if self.use_linear_projection:
																				A_ 						=		hidden_states.reshape(lowerCAmelCase   ,     height * width   ,     lowerCAmelCase  )
																				A_ 						=		self.proj_in(lowerCAmelCase  )
													else:
																				A_ 						=		self.proj_in(lowerCAmelCase  )
																				A_ 						=		hidden_states.reshape(lowerCAmelCase   ,     height * width   ,     lowerCAmelCase  )

													for transformer_block in self.transformer_blocks:
																				A_ 						=		transformer_block(lowerCAmelCase   ,     lowerCAmelCase   ,     deterministic=lowerCAmelCase  )

													if self.use_linear_projection:
																				A_ 						=		self.proj_out(lowerCAmelCase  )
																				A_ 						=		hidden_states.reshape(lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase  )
													else:
																				A_ 						=		hidden_states.reshape(lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase  )
																				A_ 						=		self.proj_out(lowerCAmelCase  )

													A_ 						=		hidden_states + residual
													return self.dropout_layer(lowerCAmelCase   ,     deterministic=lowerCAmelCase  )






class 				__lowerCAmelCase      (      nn.Module	):

						"""simple docstring"""






						_UpperCAmelCase		:					int
						_UpperCAmelCase		:					float      	=0.0
						_UpperCAmelCase		:					jnp.dtype      	=jnp.floataa






						def 							_UpperCAmelCase					(     self       :			Optional[Any]  ):
													# The second linear layer needs to be called
													# net_2 for now to match the index of the Sequential layer
													A_ 						=		FlaxGEGLU(self.dim   ,     self.dropout   ,     self.dtype  )
													A_ 						=		nn.Dense(self.dim   ,     dtype=self.dtype  )






						def __call__(     self       :			List[str]   ,     lowerCAmelCase       :			str   ,     lowerCAmelCase       :			Optional[int]=True  ):
													A_ 						=		self.net_a(lowerCAmelCase   ,     deterministic=lowerCAmelCase  )
													A_ 						=		self.net_a(lowerCAmelCase  )
													return hidden_states






class 				__lowerCAmelCase      (      nn.Module	):

						"""simple docstring"""






						_UpperCAmelCase		:					int
						_UpperCAmelCase		:					float      	=0.0
						_UpperCAmelCase		:					jnp.dtype      	=jnp.floataa






						def 							_UpperCAmelCase					(     self       :			Optional[int]  ):
													A_ 						=		self.dim * 4
													A_ 						=		nn.Dense(inner_dim * 2   ,     dtype=self.dtype  )
													A_ 						=		nn.Dropout(rate=self.dropout  )






						def __call__(     self       :			Dict   ,     lowerCAmelCase       :			Optional[Any]   ,     lowerCAmelCase       :			Tuple=True  ):
													A_ 						=		self.proj(lowerCAmelCase  )
													A_						,				A_ 						=		jnp.split(lowerCAmelCase   ,     2   ,     axis=2  )
													return self.dropout_layer(hidden_linear * nn.gelu(lowerCAmelCase  )   ,     deterministic=lowerCAmelCase  )









'''simple docstring'''
import argparse
import json
import os.path
from collections import OrderedDict

import numpy as np
import requests
import torch
from flax.training.checkpoints import restore_checkpoint
from huggingface_hub import hf_hub_download

from transformers import VivitConfig, VivitForVideoClassification, VivitImageProcessor
from transformers.image_utils import PILImageResampling







def 			a_							(  UpperCamelCase_   ):
							A_ 						=		"https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"

							with open(UpperCamelCase_	,     "wb"   ) as f:
														with requests.get(UpperCamelCase_	,     stream=UpperCamelCase_   ) as req:
																					for chunk in req.iter_content(chunk_size=2_0_4_8   ):
																												f.write(UpperCamelCase_   )







def 			a_							(  ):
							A_ 						=		VivitConfig()

							A_ 						=		4_0_0
							A_ 						=		"huggingface/label-files"
							A_ 						=		"kinetics400-id2label.json"

							A_ 						=		json.load(open(hf_hub_download(UpperCamelCase_	,     UpperCamelCase_	,     repo_type="dataset"   )	,     "r"   )   )
							A_ 						=		{int(UpperCamelCase_   ): v for k, v in idalabel.items()}
							A_ 						=		idalabel
							A_ 						=		{v: k for k, v in idalabel.items()}
							return config







def 			a_							(  ):
							A_ 						=		hf_hub_download(
							    repo_id="hf-internal-testing/spaghetti-video"	,     filename="eating_spaghetti_32_frames.npy"	,     repo_type="dataset"   )
							A_ 						=		np.load(UpperCamelCase_   )
							return list(UpperCamelCase_   )







def 			a_							(  UpperCamelCase_   ):
							if np.ndim(UpperCamelCase_   ) == 2:
														return transform_attention_bias(UpperCamelCase_   )

							elif np.ndim(UpperCamelCase_   ) == 3:
														return transform_attention_kernel(UpperCamelCase_   )

							else:
														raise Exception(f"Invalid number of dimesions: {np.ndim(UpperCamelCase_   )}"   )







def 			a_							(  UpperCamelCase_   ):
							return current.flatten()







def 			a_							(  UpperCamelCase_   ):
							return np.reshape(UpperCamelCase_	,     (current.shape[0], current.shape[1] * current.shape[2])   ).T







def 			a_							(  UpperCamelCase_   ):
							return np.reshape(UpperCamelCase_	,     (current.shape[0] * current.shape[1], current.shape[2])   ).T







def 			a_							(  UpperCamelCase_	,     UpperCamelCase_   ):
							A_ 						=		state_dict["optimizer"]["target"]["Transformer"][f"encoderblock_{i}"]

							A_ 						=		f"encoder.layer.{i}."
							A_ 						=		{
							    prefix + "intermediate.dense.bias": state["MlpBlock_0"]["Dense_0"]["bias"],
							    prefix + "intermediate.dense.weight": np.transpose(state["MlpBlock_0"]["Dense_0"]["kernel"]   ),
							    prefix + "output.dense.bias": state["MlpBlock_0"]["Dense_1"]["bias"],
							    prefix + "output.dense.weight": np.transpose(state["MlpBlock_0"]["Dense_1"]["kernel"]   ),
							    prefix + "layernorm_before.bias": state["LayerNorm_0"]["bias"],
							    prefix + "layernorm_before.weight": state["LayerNorm_0"]["scale"],
							    prefix + "layernorm_after.bias": state["LayerNorm_1"]["bias"],
							    prefix + "layernorm_after.weight": state["LayerNorm_1"]["scale"],
							    prefix
							    + "attention.attention.query.bias": transform_attention(
							        state["MultiHeadDotProductAttention_0"]["query"]["bias"]   ),
							    prefix
							    + "attention.attention.query.weight": transform_attention(
							        state["MultiHeadDotProductAttention_0"]["query"]["kernel"]   ),
							    prefix
							    + "attention.attention.key.bias": transform_attention(state["MultiHeadDotProductAttention_0"]["key"]["bias"]   ),
							    prefix
							    + "attention.attention.key.weight": transform_attention(
							        state["MultiHeadDotProductAttention_0"]["key"]["kernel"]   ),
							    prefix
							    + "attention.attention.value.bias": transform_attention(
							        state["MultiHeadDotProductAttention_0"]["value"]["bias"]   ),
							    prefix
							    + "attention.attention.value.weight": transform_attention(
							        state["MultiHeadDotProductAttention_0"]["value"]["kernel"]   ),
							    prefix + "attention.output.dense.bias": state["MultiHeadDotProductAttention_0"]["out"]["bias"],
							    prefix
							    + "attention.output.dense.weight": transform_attention_output_weight(
							        state["MultiHeadDotProductAttention_0"]["out"]["kernel"]   ),
							}

							return new_state







def 			a_							(  UpperCamelCase_   ):
							return sum([1 if "encoderblock_" in k else 0 for k in state_dict["optimizer"]["target"]["Transformer"].keys()]   )







def 			a_							(  UpperCamelCase_	,     UpperCamelCase_=False   ):
							A_ 						=		get_n_layers(UpperCamelCase_   )

							A_ 						=		OrderedDict()

							A_ 						=		state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["bias"]
							A_ 						=		state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["scale"]

							A_ 						=		np.transpose(
							    state_dict["optimizer"]["target"]["embedding"]["kernel"]	,     (4, 3, 0, 1, 2)   )
							A_ 						=		state_dict["optimizer"]["target"]["embedding"]["bias"]

							A_ 						=		state_dict["optimizer"]["target"]["cls"]
							A_ 						=		state_dict["optimizer"]["target"]["Transformer"]["posembed_input"][
							    "pos_embedding"
							]

							for i in range(UpperCamelCase_   ):
														new_state.update(transform_state_encoder_block(UpperCamelCase_	,     UpperCamelCase_   )   )

							if classification_head:
														A_ 						=		{"vivit." + k: v for k, v in new_state.items()}
														A_ 						=		np.transpose(state_dict["optimizer"]["target"]["output_projection"]["kernel"]   )
														A_ 						=		np.transpose(state_dict["optimizer"]["target"]["output_projection"]["bias"]   )

							return {k: torch.tensor(UpperCamelCase_   ) for k, v in new_state.items()}







def 			a_							(  ):
							A_ 						=		VivitImageProcessor()

							assert extractor.do_resize is True
							assert extractor.size == {"shortest_edge": 2_5_6}
							assert extractor.do_center_crop is True
							assert extractor.crop_size == {"width": 2_2_4, "height": 2_2_4}
							assert extractor.resample == PILImageResampling.BILINEAR

							# here: https://github.com/deepmind/dmvr/blob/master/dmvr/modalities.py
							# one can seen that add_image has default values for normalization_mean and normalization_std set to 0 and 1
							# which effectively means no normalization (and ViViT does not overwrite those when calling this func)
							assert extractor.do_normalize is False
							assert extractor.do_rescale is True
							assert extractor.rescale_factor == 1 / 2_5_5

							# zero-centering = True in original implementation
							assert extractor.do_zero_centering is True

							return extractor







def 			a_							(  UpperCamelCase_   ):
							A_ 						=		"checkpoint"

							if not os.path.exists(UpperCamelCase_   ):
														download_checkpoint(UpperCamelCase_   )

							A_ 						=		restore_checkpoint(UpperCamelCase_	,     UpperCamelCase_   )
							A_ 						=		transform_state(UpperCamelCase_	,     classification_head=UpperCamelCase_   )

							A_ 						=		get_vivit_config()

							assert config.image_size == 2_2_4
							assert config.num_frames == 3_2

							A_ 						=		VivitForVideoClassification(UpperCamelCase_   )
							model.load_state_dict(UpperCamelCase_   )
							model.eval()

							A_ 						=		get_processor()

							A_ 						=		prepare_video()
							A_ 						=		extractor(UpperCamelCase_	,     return_tensors="pt"   )

							A_ 						=		model(**UpperCamelCase_   )

							A_ 						=		torch.Size([1, 4_0_0]   )
							A_ 						=		torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]   )

							assert outputs.logits.shape == expected_shape
							assert torch.allclose(outputs.logits[0, :5]	,     UpperCamelCase_	,     atol=1e-4   ), outputs.logits[0, :5]

							model.save_pretrained(UpperCamelCase_   )
							extractor.save_pretrained(UpperCamelCase_   )


if __name__ == "__main__":
							__SCREAMING_SNAKE_CASE				:							Optional[Any]		       =		argparse.ArgumentParser()

							parser.add_argument('''--output_model_name''', '''-o''', type=str, help='''Output path for the converted HuggingFace model''')

							__SCREAMING_SNAKE_CASE				:							Optional[int]		       =		parser.parse_args()
							convert(args.output_model_name)





