

'''simple docstring'''



from manim import *







class 						lowerCAmelCase__   (      a					):



 """simple docstring"""



 def       UpperCAmelCase__		(		self					:       Tuple			)  -> Tuple:


      """simple docstring"""




      __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(*__SCREAMING_SNAKE_CASE			).arrange(__SCREAMING_SNAKE_CASE					,							buff=0			)
      __SCREAMING_SNAKE_CASE     =      VGroup(*__SCREAMING_SNAKE_CASE			).arrange(__SCREAMING_SNAKE_CASE					,							buff=0			)
      __SCREAMING_SNAKE_CASE     =      VGroup(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			).arrange(__SCREAMING_SNAKE_CASE					,							buff=0			)
      __SCREAMING_SNAKE_CASE     =      Text("""CPU"""					,							font_size=24			)
      __SCREAMING_SNAKE_CASE     =      Group(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			).arrange(__SCREAMING_SNAKE_CASE					,							buff=0.5					,							aligned_edge=__SCREAMING_SNAKE_CASE			)
      cpu.move_to([-2.5, -0.5, 0]			)
      self.add(__SCREAMING_SNAKE_CASE			)

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

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

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

      __SCREAMING_SNAKE_CASE     =      []
      __SCREAMING_SNAKE_CASE     =      []

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

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

      self.add(*__SCREAMING_SNAKE_CASE					,							*__SCREAMING_SNAKE_CASE			)

      __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(*__SCREAMING_SNAKE_CASE			).arrange(__SCREAMING_SNAKE_CASE					,							buff=0			)
      __SCREAMING_SNAKE_CASE     =      VGroup(*__SCREAMING_SNAKE_CASE			).arrange(__SCREAMING_SNAKE_CASE					,							buff=0			)
      __SCREAMING_SNAKE_CASE     =      VGroup(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			).arrange(__SCREAMING_SNAKE_CASE					,							buff=0			)
      __SCREAMING_SNAKE_CASE     =      Text("""Disk"""					,							font_size=24			)
      __SCREAMING_SNAKE_CASE     =      Group(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			).arrange(__SCREAMING_SNAKE_CASE					,							buff=0.5					,							aligned_edge=__SCREAMING_SNAKE_CASE			)
      disk.move_to([-4, -1.25, 0]			)
      self.add(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)

      __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(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)

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

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

      __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(__SCREAMING_SNAKE_CASE			)			)

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

      self.play(Write(__SCREAMING_SNAKE_CASE			)			)

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

      self.play(FadeOut(__SCREAMING_SNAKE_CASE			)			)

      __SCREAMING_SNAKE_CASE     =      Arrow(start=__SCREAMING_SNAKE_CASE					,							end=__SCREAMING_SNAKE_CASE					,							color=__SCREAMING_SNAKE_CASE					,							buff=0.5			)
      a.next_to(model_arr[0].get_left()					,							__SCREAMING_SNAKE_CASE					,							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(__SCREAMING_SNAKE_CASE					,							run_time=3			)			)

      __SCREAMING_SNAKE_CASE     =      {"""run_time""": 1, """fade_in""": True, """fade_out""": True, """buff""": 0.02}

      self.play(
          Write(__SCREAMING_SNAKE_CASE			)					,							Circumscribe(model_arr[0]					,							color=__SCREAMING_SNAKE_CASE					,							**__SCREAMING_SNAKE_CASE			)					,							Circumscribe(model_cpu_arr[0]					,							color=__SCREAMING_SNAKE_CASE					,							**__SCREAMING_SNAKE_CASE			)					,							Circumscribe(gpu_rect[0]					,							color=__SCREAMING_SNAKE_CASE					,							**__SCREAMING_SNAKE_CASE			)					,							)
      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					,							__SCREAMING_SNAKE_CASE					,							buff=0.2			)

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

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

           self.play(__SCREAMING_SNAKE_CASE			)

           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]					,							**__SCREAMING_SNAKE_CASE			)					,							Circumscribe(cpu_left_col_base[i]					,							**__SCREAMING_SNAKE_CASE			)					,							Circumscribe(cpu_left_col_base[i + 1]					,							color=__SCREAMING_SNAKE_CASE					,							**__SCREAMING_SNAKE_CASE			)					,							Circumscribe(gpu_rect[0]					,							color=__SCREAMING_SNAKE_CASE					,							**__SCREAMING_SNAKE_CASE			)					,							Circumscribe(model_arr[i + 1]					,							color=__SCREAMING_SNAKE_CASE					,							**__SCREAMING_SNAKE_CASE			)					,							)
                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=__SCREAMING_SNAKE_CASE					,							**__SCREAMING_SNAKE_CASE			)					,							Circumscribe(cpu_left_col_base[-1]					,							color=__SCREAMING_SNAKE_CASE					,							**__SCREAMING_SNAKE_CASE			)					,							Circumscribe(gpu_rect[0]					,							color=__SCREAMING_SNAKE_CASE					,							**__SCREAMING_SNAKE_CASE			)					,							)

                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(__SCREAMING_SNAKE_CASE			)					,							FadeOut(__SCREAMING_SNAKE_CASE					,							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(__SCREAMING_SNAKE_CASE					,							run_time=3			)					,							MoveToTarget(__SCREAMING_SNAKE_CASE			)			)

      self.wait()









'''simple docstring'''



import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
  import jax







@require_flax
class 						lowerCAmelCase__   (      a ,					unittest.TestCase					):



 """simple docstring"""



 lowerCAmelCase__  					=      FlaxAutoencoderKL
 @property
 def       UpperCAmelCase__		(		self					:       Tuple			)  -> List[str]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      4
      __SCREAMING_SNAKE_CASE     =      3
      __SCREAMING_SNAKE_CASE     =      (32, 32)

      __SCREAMING_SNAKE_CASE     =      jax.random.PRNGKey(0			)
      __SCREAMING_SNAKE_CASE     =      jax.random.uniform(__SCREAMING_SNAKE_CASE					,							((batch_size, num_channels) + sizes)			)

      return {"sample": image, "prng_key": prng_key}


 def       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> Any:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      {
          """block_out_channels""": [32, 64],
          """in_channels""": 3,
          """out_channels""": 3,
          """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
          """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
          """latent_channels""": 4,
      }
      __SCREAMING_SNAKE_CASE     =      self.dummy_input
      return init_dict, inputs_dict






