




from manim import *





class 	lowerCamelCase_							(				_SCREAMING_SNAKE_CASE ):


 '''simple docstring'''







 def 				lowerCAmelCase_  (							self  :		List[str]						):
       SCREAMING_SNAKE_CASE_     						=      Rectangle(height=0.5    ,    width=0.5						)
       SCREAMING_SNAKE_CASE_     						=      Rectangle(height=0.46    ,    width=0.46						).set_stroke(width=0						)

       SCREAMING_SNAKE_CASE_     						=      Rectangle(height=0.25    ,    width=0.25						)

       SCREAMING_SNAKE_CASE_     						=      [mem.copy() for i in range(6						)]
       SCREAMING_SNAKE_CASE_     						=      [mem.copy() for i in range(6						)]
       SCREAMING_SNAKE_CASE_     						=      VGroup(*_lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0						)
       SCREAMING_SNAKE_CASE_     						=      VGroup(*_lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0						)
       SCREAMING_SNAKE_CASE_     						=      VGroup(_lowerCAmelCase    ,    _lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0						)
       SCREAMING_SNAKE_CASE_     						=      Text('CPU'    ,    font_size=24						)
       SCREAMING_SNAKE_CASE_     						=      Group(_lowerCAmelCase    ,    _lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0.5    ,    aligned_edge=_lowerCAmelCase						)
       cpu.move_to([-2.5, -0.5, 0]						)
       self.add(_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      [mem.copy() for i in range(4						)]
       SCREAMING_SNAKE_CASE_     						=      VGroup(*_lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0						)
       SCREAMING_SNAKE_CASE_     						=      Text('GPU'    ,    font_size=24						)
       SCREAMING_SNAKE_CASE_     						=      Group(_lowerCAmelCase    ,    _lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0.5    ,    aligned_edge=_lowerCAmelCase						)
       gpu.move_to([-1, -1, 0]						)
       self.add(_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      [mem.copy() for i in range(6						)]
       SCREAMING_SNAKE_CASE_     						=      VGroup(*_lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0						)

       SCREAMING_SNAKE_CASE_     						=      Text('Model'    ,    font_size=24						)
       SCREAMING_SNAKE_CASE_     						=      Group(_lowerCAmelCase    ,    _lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0.5    ,    aligned_edge=_lowerCAmelCase						)
       model.move_to([3, -1.0, 0]						)
       self.add(_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      []
       SCREAMING_SNAKE_CASE_     						=      []

       for i, rect in enumerate(_lowerCAmelCase						):
             SCREAMING_SNAKE_CASE_     						=      fill.copy().set_fill(_lowerCAmelCase    ,    opacity=0.8						)
             target.move_to(_lowerCAmelCase						)
             model_arr.append(_lowerCAmelCase						)

             SCREAMING_SNAKE_CASE_     						=      Rectangle(height=0.46    ,    width=0.46						).set_stroke(width=0.0						).set_fill(_lowerCAmelCase    ,    opacity=0.8						)
             cpu_target.move_to(cpu_left_col_base[i]						)
             model_cpu_arr.append(_lowerCAmelCase						)

       self.add(*_lowerCAmelCase    ,    *_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      [meta_mem.copy() for i in range(6						)]
       SCREAMING_SNAKE_CASE_     						=      [meta_mem.copy() for i in range(6						)]
       SCREAMING_SNAKE_CASE_     						=      VGroup(*_lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0						)
       SCREAMING_SNAKE_CASE_     						=      VGroup(*_lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0						)
       SCREAMING_SNAKE_CASE_     						=      VGroup(_lowerCAmelCase    ,    _lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0						)
       SCREAMING_SNAKE_CASE_     						=      Text('Disk'    ,    font_size=24						)
       SCREAMING_SNAKE_CASE_     						=      Group(_lowerCAmelCase    ,    _lowerCAmelCase						).arrange(_lowerCAmelCase    ,    buff=0.5    ,    aligned_edge=_lowerCAmelCase						)
       disk.move_to([-4, -1.25, 0]						)
       self.add(_lowerCAmelCase    ,    _lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      Square(side_length=2.2						)
       key.move_to([-5, 2, 0]						)

       SCREAMING_SNAKE_CASE_     						=      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(_lowerCAmelCase    ,    _lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      MarkupText(
           F"<span fgcolor='{BLUE}'>●</span> Checkpoint"    ,    font_size=18    ,    )

       blue_text.next_to(_lowerCAmelCase    ,    DOWN * 2.4    ,    aligned_edge=key_text.get_left()						)
       self.add(_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      MarkupText(
           F"Now watch as an input is passed through the model\nand how the memory is utilized and handled."    ,    font_size=24    ,    )
       step_a.move_to([2, 2, 0]						)

       self.play(Write(_lowerCAmelCase						)						)

       SCREAMING_SNAKE_CASE_     						=      Square(0.3						)
       input.set_fill(_lowerCAmelCase    ,    opacity=1.0						)
       input.set_stroke(width=0.0						)
       input.next_to(model_base[0]    ,    _lowerCAmelCase    ,    buff=0.5						)

       self.play(Write(_lowerCAmelCase						)						)

       input.generate_target()
       input.target.next_to(model_arr[0]    ,    direction=_lowerCAmelCase    ,    buff=0.02						)
       self.play(MoveToTarget(_lowerCAmelCase						)						)

       self.play(FadeOut(_lowerCAmelCase						)						)

       SCREAMING_SNAKE_CASE_     						=      Arrow(start=_lowerCAmelCase    ,    end=_lowerCAmelCase    ,    color=_lowerCAmelCase    ,    buff=0.5						)
       a.next_to(model_arr[0].get_left()    ,    _lowerCAmelCase    ,    buff=0.2						)

       model_cpu_arr[0].generate_target()
       model_cpu_arr[0].target.move_to(gpu_rect[0]						)

       SCREAMING_SNAKE_CASE_     						=      MarkupText(
           F"As the input reaches a layer, the hook triggers\nand weights are moved from the CPU\nto the GPU and back."    ,    font_size=24    ,    )
       step_a.move_to([2, 2, 0]						)

       self.play(Write(_lowerCAmelCase    ,    run_time=3						)						)

       SCREAMING_SNAKE_CASE_     						=      {'run_time': 1, 'fade_in': True, 'fade_out': True, 'buff': 0.02}

       self.play(
           Write(_lowerCAmelCase						)    ,    Circumscribe(model_arr[0]    ,    color=_lowerCAmelCase    ,    **_lowerCAmelCase						)    ,    Circumscribe(model_cpu_arr[0]    ,    color=_lowerCAmelCase    ,    **_lowerCAmelCase						)    ,    Circumscribe(gpu_rect[0]    ,    color=_lowerCAmelCase    ,    **_lowerCAmelCase						)    ,    )
       self.play(MoveToTarget(model_cpu_arr[0]						)						)

       SCREAMING_SNAKE_CASE_     						=      a.copy()
       for i in range(6						):
             a_c.next_to(model_arr[i].get_right() + 0.02    ,    _lowerCAmelCase    ,    buff=0.2						)

             input.generate_target()
             input.target.move_to(model_arr[i].get_right() + 0.02						)

             SCREAMING_SNAKE_CASE_     						=      AnimationGroup(
                 FadeOut(_lowerCAmelCase    ,    run_time=0.5						)    ,    MoveToTarget(_lowerCAmelCase    ,    run_time=0.5						)    ,    FadeIn(_lowerCAmelCase    ,    run_time=0.5						)    ,    lag_ratio=0.2						)

             self.play(_lowerCAmelCase						)

             model_cpu_arr[i].generate_target()
             model_cpu_arr[i].target.move_to(cpu_left_col_base[i]						)

             if i < 5:
                   model_cpu_arr[i + 1].generate_target()
                   model_cpu_arr[i + 1].target.move_to(gpu_rect[0]						)
                   if i >= 1:
                         SCREAMING_SNAKE_CASE_     						=      0.7

                   self.play(
                       Circumscribe(model_arr[i]    ,    **_lowerCAmelCase						)    ,    Circumscribe(cpu_left_col_base[i]    ,    **_lowerCAmelCase						)    ,    Circumscribe(cpu_left_col_base[i + 1]    ,    color=_lowerCAmelCase    ,    **_lowerCAmelCase						)    ,    Circumscribe(gpu_rect[0]    ,    color=_lowerCAmelCase    ,    **_lowerCAmelCase						)    ,    Circumscribe(model_arr[i + 1]    ,    color=_lowerCAmelCase    ,    **_lowerCAmelCase						)    ,    )
                   if i < 1:
                         self.play(
                             MoveToTarget(model_cpu_arr[i]						)    ,    MoveToTarget(model_cpu_arr[i + 1]						)    ,    )
                   else:
                         self.play(
                             MoveToTarget(model_cpu_arr[i]    ,    run_time=0.7						)    ,    MoveToTarget(model_cpu_arr[i + 1]    ,    run_time=0.7						)    ,    )
             else:
                   model_cpu_arr[i].generate_target()
                   model_cpu_arr[i].target.move_to(cpu_left_col_base[-1]						)
                   input.generate_target()
                   input.target.next_to(model_arr[-1].get_right()    ,    RIGHT + 0.02    ,    buff=0.2						)

                   self.play(
                       Circumscribe(model_arr[-1]    ,    color=_lowerCAmelCase    ,    **_lowerCAmelCase						)    ,    Circumscribe(cpu_left_col_base[-1]    ,    color=_lowerCAmelCase    ,    **_lowerCAmelCase						)    ,    Circumscribe(gpu_rect[0]    ,    color=_lowerCAmelCase    ,    **_lowerCAmelCase						)    ,    )

                   self.play(MoveToTarget(model_cpu_arr[i]						)						)

             SCREAMING_SNAKE_CASE_     						=      a_c
             SCREAMING_SNAKE_CASE_     						=      a_c.copy()

       input.generate_target()
       input.target.next_to(model_base[-1]    ,    RIGHT + 0.02    ,    buff=0.5						)
       self.play(
           FadeOut(_lowerCAmelCase						)    ,    FadeOut(_lowerCAmelCase    ,    run_time=0.5						)    ,    )

       SCREAMING_SNAKE_CASE_     						=      MarkupText(F"Inference on a model too large for GPU memory\nis successfully completed."    ,    font_size=24						)
       step_a.move_to([2, 2, 0]						)

       self.play(Write(_lowerCAmelCase    ,    run_time=3						)    ,    MoveToTarget(_lowerCAmelCase						)						)

       self.wait()





from packaging import version

from .import_utils import is_accelerate_available


if is_accelerate_available():
   import accelerate

def  UpperCAmelCase_							(				__UpperCAmelCase						:     Optional[int]							) ->			int:
      if not is_accelerate_available():
            return method
      SCREAMING_SNAKE_CASE_     						=      version.parse(accelerate.__version__							).base_version
      if version.parse(__UpperCAmelCase							) < version.parse('0.17.0'							):
            return method

      def wrapper(self						:     Optional[int]      ,      *__UpperCAmelCase						:     Optional[Any]      ,      **__UpperCAmelCase						:     Optional[Any]							):
            if hasattr(self      ,      '_hf_hook'							) and hasattr(self._hf_hook      ,      'pre_forward'							):
                  self._hf_hook.pre_forward(self							)
            return method(self      ,      *__UpperCAmelCase      ,      **__UpperCAmelCase							)

      return wrapper