code: 
  
  
  import json
  import os
  from collections.abc import Mapping
  from typing import Dict, List, Optional, Union
  
  import numpy as np
  import torch
  
  from .imports import is_safetensors_available
  def     UpperCAmelCase__		(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_=None       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       snake_case					=       None
       # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
       if str(weight.dtype       ) == "torch.bfloat16":
            # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
            snake_case					=       weight.view(torch.intaa       )
            snake_case					=       '''bfloat16'''
       snake_case					=       weight.cpu().numpy()
       snake_case					=       os.path.join(UpperCamelCase_							,F'''{weight_name}.dat'''       )
       if index is not None:
            if dtype is None:
                 snake_case					=       str(array.dtype       )
            snake_case					=       {'''dtype''': dtype, '''shape''': list(array.shape       )}
       if array.ndim == 0:
            snake_case					=       array[None]
       snake_case					=       np.memmap(UpperCamelCase_							,dtype=array.dtype							,mode='''w+'''							,shape=array.shape       )
       snake_case					=       array[:]
       file_array.flush()
       return index
  def     UpperCAmelCase__		(UpperCamelCase_							,UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       snake_case					=       tuple(weight_info['''shape''']       )
       if shape == ():
            # NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
            snake_case					=       (1,)
  
       snake_case					=       weight_info['''dtype''']
       if dtype == "bfloat16":
            # NumPy does not support bfloat16 so this was saved as a int16
            snake_case					=       '''int16'''
  
       snake_case					=       np.memmap(UpperCamelCase_							,dtype=UpperCamelCase_							,shape=UpperCamelCase_							,mode='''r'''       )
  
       if len(weight_info['''shape''']       ) == 0:
            snake_case					=       weight[0]
       snake_case					=       torch.tensor(UpperCamelCase_       )
       if weight_info["dtype"] == "bfloat16":
            snake_case					=       weight.view(torch.bfloataa       )
  
       return weight
  def     UpperCAmelCase__		(UpperCamelCase_							,UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       if index is None or len(UpperCamelCase_       ) == 0:
            # Nothing to save
            return
  
       snake_case					=       os.path.join(UpperCamelCase_							,'''index.json'''       )
       if os.path.isfile(UpperCamelCase_       ):
            with open(UpperCamelCase_							,'''r'''							,encoding='''utf-8'''       ) as f:
                 snake_case					=       json.load(UpperCamelCase_       )
       else:
            snake_case					=       {}
       current_index.update(UpperCamelCase_       )
  
       with open(UpperCamelCase_							,'''w'''							,encoding='''utf-8'''       ) as f:
            json.dump(UpperCamelCase_							,UpperCamelCase_							,indent=2       )
  def     UpperCAmelCase__		(UpperCamelCase_							,UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       os.makedirs(UpperCamelCase_							,exist_ok=UpperCamelCase_       )
       snake_case					=       {}
       for name, parameter in state_dict.items():
            snake_case					=       offload_weight(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_							,index=UpperCamelCase_       )
  
       # Update index
       save_offload_index(UpperCamelCase_							,UpperCamelCase_       )
  
  class 						A__      (					snake_case__			):
  
  
  
  
  
  
  
      """simple docstring"""
  
      def __init__( self    ,      __snake_case    ,      __snake_case					):
           snake_case					=       dataset
           snake_case					=       prefix
      def __getitem__( self    ,      __snake_case					):
           return self.dataset[F'''{self.prefix}{key}''']
      def __iter__( self					):
           return iter([key for key in self.dataset if key.startswith(self.prefix					)]					)
  
  
  
  
      def __len__( self					):
           return len(self.dataset					)
  
  
  
  
  
  
  class 						A__      (					snake_case__			):
  
  
  
  
  
  
  
      """simple docstring"""
  
      def __init__( self    ,      __snake_case = None    ,      __snake_case = None    ,      __snake_case = None    ,      __snake_case=None    ,      ):
           if state_dict is None and save_folder is None:
                raise ValueError('''Need either a `state_dict` or a `save_folder` containing offloaded weights.'''					)
  
           snake_case					=       {} if state_dict is None else state_dict
           snake_case					=       save_folder
           if index is None and save_folder is not None:
                with open(os.path.join(__snake_case    ,      '''index.json'''					)					) as f:
                     snake_case					=       json.load(__snake_case					)
           snake_case					=       {} if index is None else index
           snake_case					=       list(self.state_dict.keys()					)
           self.all_keys.extend([key for key in self.index if key not in self.all_keys]					)
           snake_case					=       device
      def __getitem__( self    ,      __snake_case					):
           # State dict gets priority
           if key in self.state_dict:
                return self.state_dict[key]
           snake_case					=       self.index[key]
           if weight_info.get('''safetensors_file'''					) is not None:
                if not is_safetensors_available():
                     raise ImportError('''These offloaded weights require the use of safetensors: `pip install safetensors`.'''					)
  
                from safetensors import safe_open
  
                snake_case					=       '''cpu''' if self.device is None else self.device
                with safe_open(weight_info['''safetensors_file''']    ,      framework='''pt'''    ,      device=__snake_case					) as f:
                     snake_case					=       f.get_tensor(weight_info.get('''weight_name'''    ,      __snake_case					)					)
  
                if "dtype" in weight_info:
                     return tensor.to(getattr(__snake_case    ,      weight_info['''dtype''']					)					)
                else:
                     return tensor
  
           snake_case					=       os.path.join(self.save_folder    ,      F'''{key}.dat'''					)
           return load_offloaded_weight(__snake_case    ,      __snake_case					)
      def __iter__( self					):
           return iter(self.all_keys					)
  
  
  
  
      def __len__( self					):
           return len(self.all_keys					)
  
  
  
  
  
  
  def     UpperCAmelCase__		(UpperCamelCase_							,UpperCamelCase_       ):
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       snake_case					=       {}
       for module_name in submodule_names:
            # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
            # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
            result.update(
                {
                    key: param
                    for key, param in state_dict.items()
                    if key == module_name or key.startswith(module_name + '''.'''       )
                }       )
       return result
  
  
  
  
  
  
  
style_context: 
  
  
  from manim import *
  
  class 						A__      (					snake_case__			):
  
  
  
  
  
  
  
      """simple docstring"""
  
      def 	a_    ( self					):
           snake_case					=       Rectangle(height=0.5    ,      width=0.5					)
           snake_case					=       Rectangle(height=0.46    ,      width=0.46					).set_stroke(width=0					)
           snake_case					=       Rectangle(height=0.25    ,      width=0.25					)
  
           snake_case					=       [mem.copy() for i in range(6					)]
           snake_case					=       [mem.copy() for i in range(6					)]
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       VGroup(__snake_case    ,      __snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       Text('''CPU'''    ,      font_size=2_4					)
           snake_case					=       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					)
  
           snake_case					=       [mem.copy() for i in range(4					)]
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       Text('''GPU'''    ,      font_size=2_4					)
           snake_case					=       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					)
  
           snake_case					=       [mem.copy() for i in range(6					)]
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
  
           snake_case					=       Text('''Model'''    ,      font_size=2_4					)
           snake_case					=       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					)
  
           snake_case					=       []
           snake_case					=       []
  
           for i, rect in enumerate(__snake_case					):
                rect.set_stroke(__snake_case					)
  
                snake_case					=       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					)
  
           snake_case					=       [meta_mem.copy() for i in range(6					)]
           snake_case					=       [meta_mem.copy() for i in range(6					)]
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       VGroup(__snake_case    ,      __snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       Text('''Disk'''    ,      font_size=2_4					)
           snake_case					=       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					)
  
           snake_case					=       []
  
           for i in range(6					):
                snake_case					=       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					)
  
           snake_case					=       Square(side_length=2.2					)
           key.move_to([-5, 2, 0]					)
  
           snake_case					=       MarkupText(
               F'''<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'''    ,      font_size=1_8    ,      )
  
           key_text.move_to([-5, 2.4, 0]					)
  
           self.add(__snake_case    ,      __snake_case					)
  
           snake_case					=       MarkupText(
               F'''<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'''    ,      font_size=1_8    ,      )
  
           blue_text.next_to(__snake_case    ,      DOWN * 2.4    ,      aligned_edge=key_text.get_left()					)
           self.add(__snake_case					)
  
           snake_case					=       MarkupText(F'''The offloaded weights are all sent to the CPU.'''    ,      font_size=2_4					)
           step_a.move_to([2, 2, 0]					)
  
           self.play(Write(__snake_case    ,      run_time=3					)					)
  
           for i in range(6					):
                snake_case					=       cpu_disk_arr[i]
                snake_case					=       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					)					)
  
           snake_case					=       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=2_4    ,      )
           step_a.move_to([2, 2, 0]					)
  
           self.play(Write(__snake_case    ,      run_time=3					)					)
  
           snake_case					=       []
           snake_case					=       []
           for i in range(6					):
                snake_case					=       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()
  
  
  
  
  