


from manim import *



class  UpperCAmelCase_						(				_a):
       '''simple docstring'''






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Union[str, Any] 						=		Rectangle(height=0.5			,       width=0.5    )
              UpperCamelCase     :						Optional[int] 						=		Rectangle(height=0.46			,       width=0.46    ).set_stroke(width=0    )
              UpperCamelCase     :						Optional[Any] 						=		Rectangle(height=0.25			,       width=0.25    )

              UpperCamelCase     :						List[Any] 						=		[mem.copy() for i in range(6    )]
              UpperCamelCase     :						List[Any] 						=		[mem.copy() for i in range(6    )]
              UpperCamelCase     :						int 						=		VGroup(*__SCREAMING_SNAKE_CASE    ).arrange(__SCREAMING_SNAKE_CASE			,       buff=0    )
              UpperCamelCase     :						str 						=		VGroup(*__SCREAMING_SNAKE_CASE    ).arrange(__SCREAMING_SNAKE_CASE			,       buff=0    )
              UpperCamelCase     :						Optional[Any] 						=		VGroup(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ).arrange(__SCREAMING_SNAKE_CASE			,       buff=0    )
              UpperCamelCase     :						Tuple 						=		Text('''CPU'''			,       font_size=24    )
              UpperCamelCase     :						int 						=		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    )

              UpperCamelCase     :						Any 						=		[mem.copy() for i in range(4    )]
              UpperCamelCase     :						str 						=		VGroup(*__SCREAMING_SNAKE_CASE    ).arrange(__SCREAMING_SNAKE_CASE			,       buff=0    )
              UpperCamelCase     :						List[str] 						=		Text('''GPU'''			,       font_size=24    )
              UpperCamelCase     :						int 						=		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    )

              UpperCamelCase     :						Dict 						=		[mem.copy() for i in range(6    )]
              UpperCamelCase     :						Union[str, Any] 						=		VGroup(*__SCREAMING_SNAKE_CASE    ).arrange(__SCREAMING_SNAKE_CASE			,       buff=0    )

              UpperCamelCase     :						Optional[int] 						=		Text('''Model'''			,       font_size=24    )
              UpperCamelCase     :						int 						=		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    )

              UpperCamelCase     :						Optional[Any] 						=		[]
              UpperCamelCase     :						Dict 						=		[]

              for i, rect in enumerate(__SCREAMING_SNAKE_CASE    ):
                     rect.set_stroke(__SCREAMING_SNAKE_CASE    )

                     UpperCamelCase     :						Tuple 						=		Rectangle(height=0.46 / 4			,       width=0.46 / 3    ).set_stroke(width=0.0    ).set_fill(__SCREAMING_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=__SCREAMING_SNAKE_CASE    )
                            cpu_target.set_x(cpu_target.get_x() + 0.1    )
                     elif i == 3:
                            cpu_target.next_to(model_cpu_arr[0]			,       direction=__SCREAMING_SNAKE_CASE			,       buff=0.0    )
                     else:
                            cpu_target.next_to(model_cpu_arr[i - 1]			,       direction=__SCREAMING_SNAKE_CASE			,       buff=0.0    )
                     self.add(__SCREAMING_SNAKE_CASE    )
                     model_cpu_arr.append(__SCREAMING_SNAKE_CASE    )

              self.add(*__SCREAMING_SNAKE_CASE			,       *__SCREAMING_SNAKE_CASE    )

              UpperCamelCase     :						str 						=		[meta_mem.copy() for i in range(6    )]
              UpperCamelCase     :						Optional[int] 						=		[meta_mem.copy() for i in range(6    )]
              UpperCamelCase     :						Optional[Any] 						=		VGroup(*__SCREAMING_SNAKE_CASE    ).arrange(__SCREAMING_SNAKE_CASE			,       buff=0    )
              UpperCamelCase     :						Dict 						=		VGroup(*__SCREAMING_SNAKE_CASE    ).arrange(__SCREAMING_SNAKE_CASE			,       buff=0    )
              UpperCamelCase     :						str 						=		VGroup(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ).arrange(__SCREAMING_SNAKE_CASE			,       buff=0    )
              UpperCamelCase     :						Optional[int] 						=		Text('''Disk'''			,       font_size=24    )
              UpperCamelCase     :						Optional[Any] 						=		Group(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    ).arrange(__SCREAMING_SNAKE_CASE			,       buff=0.5			,       aligned_edge=__SCREAMING_SNAKE_CASE    )
              disk.move_to([-4.0, -1.25, 0]    )
              self.add(__SCREAMING_SNAKE_CASE			,       __SCREAMING_SNAKE_CASE    )

              UpperCamelCase     :						Optional[Any] 						=		[]

              for i in range(6    ):
                     UpperCamelCase     :						Optional[Any] 						=		fill.copy().set_fill(__SCREAMING_SNAKE_CASE			,       opacity=0.8    )
                     target.move_to(disk_left_col_base[i]    ).scale(0.5    )
                     cpu_disk_arr.append(__SCREAMING_SNAKE_CASE    )

              self.add(*__SCREAMING_SNAKE_CASE    )

              UpperCamelCase     :						Union[str, Any] 						=		Square(side_length=2.2    )
              key.move_to([-5, 2, 0]    )

              UpperCamelCase     :						int 						=		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    )

              UpperCamelCase     :						Optional[int] 						=		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    )

              UpperCamelCase     :						List[str] 						=		MarkupText(f"""The offloaded weights are all sent to the CPU."""			,       font_size=24    )
              step_a.move_to([2, 2, 0]    )

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

              for i in range(6    ):
                     UpperCamelCase     :						Optional[int] 						=		cpu_disk_arr[i]
                     UpperCamelCase     :						List[Any] 						=		rect.copy().set_fill(__SCREAMING_SNAKE_CASE			,       opacity=0.8    ).scale(2.0    )
                     cpa.generate_target()
                     cpa.target.move_to(model_base[i]    )

                     if i == 0:
                            rect.set_fill(__SCREAMING_SNAKE_CASE			,       opacity=0.8    )
                            rect.generate_target()
                            rect.target.move_to(cpu_left_col_base[0]    ).scale(2.0    )

                            self.remove(
                                *__SCREAMING_SNAKE_CASE			,       *__SCREAMING_SNAKE_CASE			,       )

                     else:
                            rect.generate_target()
                            rect.target.move_to(cpu_left_col_base[i]    ).scale(2.0    )
                     self.play(MoveToTarget(__SCREAMING_SNAKE_CASE    )			,       MoveToTarget(__SCREAMING_SNAKE_CASE    )			,       model_base[i].animate.set_stroke(__SCREAMING_SNAKE_CASE    )    )
              self.play(FadeOut(__SCREAMING_SNAKE_CASE    )    )

              UpperCamelCase     :						Dict 						=		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(__SCREAMING_SNAKE_CASE			,       run_time=3    )    )

              UpperCamelCase     :						Optional[int] 						=		[]
              UpperCamelCase     :						Any 						=		[]
              for i in range(6    ):
                     UpperCamelCase     :						List[Any] 						=		Arrow(start=__SCREAMING_SNAKE_CASE			,       end=__SCREAMING_SNAKE_CASE			,       color=__SCREAMING_SNAKE_CASE			,       buff=0.5    )
                     a.next_to(model_base[i].get_left()			,       __SCREAMING_SNAKE_CASE			,       buff=0.2    )
                     arrows.append(__SCREAMING_SNAKE_CASE    )
                     animations.append(Write(__SCREAMING_SNAKE_CASE    )    )
              self.play(*__SCREAMING_SNAKE_CASE    )
              self.wait()





import os

from apache_beam.io.filesystems import FileSystems
from apache_beam.pipeline import Pipeline

from .logging import get_logger


__UpperCAmelCase			:			Optional[Any]			      =					2 << 20  # 2mb
__UpperCAmelCase			:			List[Any]			      =					get_logger(__name__)



class  UpperCAmelCase_						(				_a):
       '''simple docstring'''






       def 			_lowercase (  self    ):




              """simple docstring"""




              UpperCamelCase     :						Optional[Any] 						=		self._options.get_all_options().get('''runner'''    )
              return runner in [None, "DirectRunner", "PortableRunner"]




def   a						(				SCREAMING_SNAKE_CASE_   :		str				,      SCREAMING_SNAKE_CASE_   :		int				,      SCREAMING_SNAKE_CASE_   :		Optional[Any]=False							):






       """simple docstring"""





       UpperCamelCase     :						int 						=		FileSystems
       if fs.exists(SCREAMING_SNAKE_CASE_							):
              if force_upload:
                     logger.info(F"""Remote path already exist: {remote_file_path}. Overwriting it as force_upload=True."""							)
              else:
                     logger.info(F"""Remote path already exist: {remote_file_path}. Skipping it as force_upload=False."""							)
                     return
       with fs.create(SCREAMING_SNAKE_CASE_							) as remote_file:
              with open(SCREAMING_SNAKE_CASE_				,      '''rb'''							) as local_file:
                     UpperCamelCase     :						List[str] 						=		local_file.read(SCREAMING_SNAKE_CASE_							)
                     while chunk:
                            remote_file.write(SCREAMING_SNAKE_CASE_							)
                            UpperCamelCase     :						Any 						=		local_file.read(SCREAMING_SNAKE_CASE_							)


def   a						(				SCREAMING_SNAKE_CASE_   :		str				,      SCREAMING_SNAKE_CASE_   :		Optional[Any]				,      SCREAMING_SNAKE_CASE_   :		List[str]=False							):






       """simple docstring"""





       UpperCamelCase     :						List[Any] 						=		FileSystems
       if os.path.exists(SCREAMING_SNAKE_CASE_							):
              if force_download:
                     logger.info(F"""Local path already exist: {remote_file_path}. Overwriting it as force_upload=True."""							)
              else:
                     logger.info(F"""Local path already exist: {remote_file_path}. Skipping it as force_upload=False."""							)
                     return
       with fs.open(SCREAMING_SNAKE_CASE_							) as remote_file:
              with open(SCREAMING_SNAKE_CASE_				,      '''wb'''							) as local_file:
                     UpperCamelCase     :						List[str] 						=		remote_file.read(SCREAMING_SNAKE_CASE_							)
                     while chunk:
                            local_file.write(SCREAMING_SNAKE_CASE_							)
                            UpperCamelCase     :						str 						=		remote_file.read(SCREAMING_SNAKE_CASE_							)

