

"""simple docstring"""






from manim import *
class        _lowercase					(     __UpperCAmelCase     ):







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      Rectangle(height=0.5 ,						width=0.5       )
      __magic_name__     =      Rectangle(height=0.4_6 ,						width=0.4_6       ).set_stroke(width=0       )

      __magic_name__     =      Rectangle(height=0.2_5 ,						width=0.2_5       )

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

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

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

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

      __magic_name__     =      []
      __magic_name__     =      []

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

       __magic_name__     =      Rectangle(height=0.4_6 ,						width=0.4_6       ).set_stroke(width=0.0       ).set_fill(UpperCamelCase_ ,						opacity=0.8       )
       cpu_target.move_to(cpu_left_col_base[i]       )
       model_cpu_arr.append(UpperCamelCase_       )

      self.add(*UpperCamelCase_ ,						*UpperCamelCase_       )

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

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

      __magic_name__     =      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(UpperCamelCase_ ,						UpperCamelCase_       )

      __magic_name__     =      MarkupText(
          f'''<span fgcolor=\'{BLUE}\'>●</span> Checkpoint''' ,						font_size=18 ,						)

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

      __magic_name__     =      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(UpperCamelCase_       )       )

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

      self.play(Write(UpperCamelCase_       )       )

      input.generate_target()
      input.target.next_to(model_arr[0] ,						direction=UpperCamelCase_ ,						buff=0.0_2       )
      self.play(MoveToTarget(UpperCamelCase_       )       )

      self.play(FadeOut(UpperCamelCase_       )       )

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

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

      __magic_name__     =      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(UpperCamelCase_ ,						run_time=3       )       )

      __magic_name__     =      {'''run_time''': 1, '''fade_in''': True, '''fade_out''': True, '''buff''': 0.0_2}

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

      __magic_name__     =      a.copy()
      for i in range(6       ):
       a_c.next_to(model_arr[i].get_right() + 0.0_2 ,						UpperCamelCase_ ,						buff=0.2       )

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

       __magic_name__     =      AnimationGroup(
           FadeOut(UpperCamelCase_ ,						run_time=0.5       ) ,						MoveToTarget(UpperCamelCase_ ,						run_time=0.5       ) ,						FadeIn(UpperCamelCase_ ,						run_time=0.5       ) ,						lag_ratio=0.2       )

       self.play(UpperCamelCase_       )

       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:
         __magic_name__     =      0.7

        self.play(
            Circumscribe(model_arr[i] ,						**UpperCamelCase_       ) ,						Circumscribe(cpu_left_col_base[i] ,						**UpperCamelCase_       ) ,						Circumscribe(cpu_left_col_base[i + 1] ,						color=UpperCamelCase_ ,						**UpperCamelCase_       ) ,						Circumscribe(gpu_rect[0] ,						color=UpperCamelCase_ ,						**UpperCamelCase_       ) ,						Circumscribe(model_arr[i + 1] ,						color=UpperCamelCase_ ,						**UpperCamelCase_       ) ,						)
        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.0_2 ,						buff=0.2       )

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

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

       __magic_name__     =      a_c
       __magic_name__     =      a_c.copy()

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

      __magic_name__     =      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(UpperCamelCase_ ,						run_time=3       ) ,						MoveToTarget(UpperCamelCase_       )       )

      self.wait()




"""simple docstring"""






import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
 from transformers import AutoModelForSeqaSeqLM, AutoTokenizer
@require_torch
@require_sentencepiece
@require_tokenizers
class        _lowercase					(     unittest.TestCase     ):







     @slow
     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      AutoModelForSeqaSeqLM.from_pretrained('''google/mt5-small''' ,						return_dict=UpperCamelCase_       ).to(UpperCamelCase_       )
      __magic_name__     =      AutoTokenizer.from_pretrained('''google/mt5-small'''       )

      __magic_name__     =      tokenizer('''Hello there''' ,						return_tensors='''pt'''       ).input_ids
      __magic_name__     =      tokenizer('''Hi I am''' ,						return_tensors='''pt'''       ).input_ids

      __magic_name__     =      model(input_ids.to(UpperCamelCase_       ) ,						labels=labels.to(UpperCamelCase_       )       ).loss
      __magic_name__     =      -(labels.shape[-1] * loss.item())

      __magic_name__     =      -8_4.9_1_2_7
      self.assertTrue(abs(mtf_score - EXPECTED_SCORE       ) < 1E-4       )

