code: 
  
  
  
  
  """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 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Any:
   A__ =       "https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"
  
   with open(lowercase_    ,				"wb"		) as f:
    with requests.get(lowercase_    ,				stream=lowercase_		) as req:
     for chunk in req.iter_content(chunk_size=20_48		):
      f.write(lowercase_		)
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							)		->	VivitConfig:
   A__ =       VivitConfig()
  
   A__ =       4_00
   A__ =       "huggingface/label-files"
   A__ =       "kinetics400-id2label.json"
  
   A__ =       json.load(open(hf_hub_download(lowercase_    ,				lowercase_    ,				repo_type="dataset"		)    ,				"r"		)		)
   A__ =       {int(lowercase_		): v for k, v in idalabel.items()}
   A__ =       idalabel
   A__ =       {v: k for k, v in idalabel.items()}
   return config
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							)		->	Optional[int]:
   A__ =       hf_hub_download(
       repo_id="hf-internal-testing/spaghetti-video"    ,				filename="eating_spaghetti_32_frames.npy"    ,				repo_type="dataset"		)
   A__ =       np.load(lowercase_		)
   return list(lowercase_		)
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Tuple:
   if np.ndim(lowercase_		) == 2:
    return transform_attention_bias(lowercase_		)
  
   elif np.ndim(lowercase_		) == 3:
    return transform_attention_kernel(lowercase_		)
  
   else:
    raise Exception(f"""Invalid number of dimesions: {np.ndim(lowercase_		)}"""		)
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Any:
   return current.flatten()
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	str:
   return np.reshape(lowercase_    ,				(current.shape[0], current.shape[1] * current.shape[2])		).T
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	List[str]:
   return np.reshape(lowercase_    ,				(current.shape[0] * current.shape[1], current.shape[2])		).T
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							lowercase_    ,				lowercase_		)		->	int:
   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 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Any:
   return sum([1 if "encoderblock_" in k else 0 for k in state_dict["optimizer"]["target"]["Transformer"].keys()]		)
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							lowercase_    ,				lowercase_=False		)		->	Union[str, Any]:
   A__ =       get_n_layers(lowercase_		)
  
   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(lowercase_		):
    new_state.update(transform_state_encoder_block(lowercase_    ,				lowercase_		)		)
  
   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(lowercase_		) for k, v in new_state.items()}
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							)		->	VivitImageProcessor:
   A__ =       VivitImageProcessor()
  
   assert extractor.do_resize is True
   assert extractor.size == {"shortest_edge": 2_56}
   assert extractor.do_center_crop is True
   assert extractor.crop_size == {"width": 2_24, "height": 2_24}
   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_55
  
   # zero-centering = True in original implementation
   assert extractor.do_zero_centering is True
  
   return extractor
  
  
  
  
  
  
  
  def 						_SCREAMING_SNAKE_CASE     (							lowercase_		)		->	Optional[int]:
   A__ =       "checkpoint"
  
   if not os.path.exists(lowercase_		):
    download_checkpoint(lowercase_		)
  
   A__ =       restore_checkpoint(lowercase_    ,				lowercase_		)
   A__ =       transform_state(lowercase_    ,				classification_head=lowercase_		)
  
   A__ =       get_vivit_config()
  
   assert config.image_size == 2_24
   assert config.num_frames == 32
  
   A__ =       VivitForVideoClassification(lowercase_		)
   model.load_state_dict(lowercase_		)
   model.eval()
  
   A__ =       get_processor()
  
   A__ =       prepare_video()
   A__ =       extractor(lowercase_    ,				return_tensors="pt"		)
  
   A__ =       model(**lowercase_		)
  
   A__ =       torch.Size([1, 4_00]		)
   A__ =       torch.tensor([-1.0_5_4_3, 2.0_7_6_4, -0.2_1_0_4, 0.4_4_3_9, -0.9_6_5_8]		)
  
   assert outputs.logits.shape == expected_shape
   assert torch.allclose(outputs.logits[0, :5]    ,				lowercase_    ,				atol=1E-4		), outputs.logits[0, :5]
  
   model.save_pretrained(lowercase_		)
   extractor.save_pretrained(lowercase_		)
  
  
  if __name__ == "__main__":
        SCREAMING_SNAKE_CASE       = argparse.ArgumentParser()
  
        parser.add_argument("--output_model_name", "-o", type=str, help="Output path for the converted HuggingFace model")
  
        SCREAMING_SNAKE_CASE       = parser.parse_args()
        convert(args.output_model_name)
  
  
  
  
  
  
style_context: 
  
  
  
  
  """simple docstring"""
  
  
  
  from manim import *
  
  
  
  
  class 			UpperCAmelCase_						(		A_							):
  
  
  
  
      def 							__magic_name__							(					self		: Any   )   ->							Optional[int]:
  
  
  
  
       '''simple docstring'''
  
       A__ =       Rectangle(height=0.5			,				width=0.5   )
       A__ =       Rectangle(height=0.46			,				width=0.46   ).set_stroke(width=0   )
       A__ =       Rectangle(height=0.25			,				width=0.25   )
  
       A__ =       [mem.copy() for i in range(6   )]
       A__ =       [mem.copy() for i in range(6   )]
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       VGroup(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       Text("CPU"			,				font_size=24   )
       A__ =       Group(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0.5			,				aligned_edge=snake_case_   )
       cpu.move_to([-2.5, -0.5, 0]   )
       self.add(snake_case_   )
  
       A__ =       [mem.copy() for i in range(4   )]
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       Text("GPU"			,				font_size=24   )
       A__ =       Group(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0.5			,				aligned_edge=snake_case_   )
       gpu.move_to([-1, -1, 0]   )
       self.add(snake_case_   )
  
       A__ =       [mem.copy() for i in range(6   )]
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
  
       A__ =       Text("Model"			,				font_size=24   )
       A__ =       Group(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0.5			,				aligned_edge=snake_case_   )
       model.move_to([3, -1.0, 0]   )
       self.add(snake_case_   )
  
       A__ =       []
       A__ =       []
  
       for i, rect in enumerate(snake_case_   ):
        rect.set_stroke(snake_case_   )
  
        A__ =       Rectangle(height=0.46 / 4			,				width=0.46 / 3   ).set_stroke(width=0.0   ).set_fill(snake_case_			,				opacity=0.7   )
  
        if i == 0:
         cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT   )			,				buff=0.02			,				direction=snake_case_   )
         cpu_target.set_x(cpu_target.get_x() + 0.1   )
        elif i == 3:
         cpu_target.next_to(model_cpu_arr[0]			,				direction=snake_case_			,				buff=0.0   )
        else:
         cpu_target.next_to(model_cpu_arr[i - 1]			,				direction=snake_case_			,				buff=0.0   )
        self.add(snake_case_   )
        model_cpu_arr.append(snake_case_   )
  
       self.add(*snake_case_			,				*snake_case_   )
  
       A__ =       [meta_mem.copy() for i in range(6   )]
       A__ =       [meta_mem.copy() for i in range(6   )]
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       VGroup(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       Text("Disk"			,				font_size=24   )
       A__ =       Group(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0.5			,				aligned_edge=snake_case_   )
       disk.move_to([-4.0, -1.25, 0]   )
       self.add(snake_case_			,				snake_case_   )
  
       A__ =       []
  
       for i in range(6   ):
        A__ =       fill.copy().set_fill(snake_case_			,				opacity=0.8   )
        target.move_to(disk_left_col_base[i]   ).scale(0.5   )
        cpu_disk_arr.append(snake_case_   )
  
       self.add(*snake_case_   )
  
       A__ =       Square(side_length=2.2   )
       key.move_to([-5, 2, 0]   )
  
       A__ =       MarkupText(
           F"""<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"""			,				font_size=18			,				)
  
       key_text.move_to([-5, 2.4, 0]   )
  
       self.add(snake_case_			,				snake_case_   )
  
       A__ =       MarkupText(
           F"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""			,				font_size=18			,				)
  
       blue_text.next_to(snake_case_			,				DOWN * 2.4			,				aligned_edge=key_text.get_left()   )
       self.add(snake_case_   )
  
       A__ =       MarkupText(F"""The offloaded weights are all sent to the CPU."""			,				font_size=24   )
       step_a.move_to([2, 2, 0]   )
  
       self.play(Write(snake_case_			,				run_time=3   )   )
  
       for i in range(6   ):
        A__ =       cpu_disk_arr[i]
        A__ =       rect.copy().set_fill(snake_case_			,				opacity=0.8   ).scale(2.0   )
        cpa.generate_target()
        cpa.target.move_to(model_base[i]   )
  
        if i == 0:
         rect.set_fill(snake_case_			,				opacity=0.8   )
         rect.generate_target()
         rect.target.move_to(cpu_left_col_base[0]   ).scale(2.0   )
  
         self.remove(
             *snake_case_			,				*snake_case_			,				)
  
        else:
         rect.generate_target()
         rect.target.move_to(cpu_left_col_base[i]   ).scale(2.0   )
        self.play(MoveToTarget(snake_case_   )			,				MoveToTarget(snake_case_   )			,				model_base[i].animate.set_stroke(snake_case_   )   )
       self.play(FadeOut(snake_case_   )   )
  
       A__ =       MarkupText(
           F"""Finally, hooks are added to each weight in the model\nto transfer the weights from CPU to GPU\n\t\tand back when needed."""			,				font_size=24			,				)
       step_a.move_to([2, 2, 0]   )
  
       self.play(Write(snake_case_			,				run_time=3   )   )
  
       A__ =       []
       A__ =       []
       for i in range(6   ):
        A__ =       Arrow(start=snake_case_			,				end=snake_case_			,				color=snake_case_			,				buff=0.5   )
        a.next_to(model_base[i].get_left()			,				snake_case_			,				buff=0.2   )
        arrows.append(snake_case_   )
        animations.append(Write(snake_case_   )   )
       self.play(*snake_case_   )
       self.wait()
  
  
  
  
  
  