




"""simple docstring"""






from manim import *




class   a__			(       UpperCamelCase_  ):





						def       __UpperCamelCase     (							self      :	List[Any])						->       str:




											"""simple docstring"""







											_lowerCAmelCase:Optional[int]									=  Rectangle(height=0.5     ,width=0.5)
											_lowerCAmelCase:Any									=  Rectangle(height=0.46     ,width=0.46).set_stroke(width=0)

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

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

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

											_lowerCAmelCase:Tuple									=  Text('''Model'''     ,font_size=24)
											_lowerCAmelCase:Union[str, Any]									=  Group(a__     ,a__).arrange(a__     ,buff=0.5     ,aligned_edge=a__)
											model.move_to([3, -1.0, 0])
											self.add(a__)

											_lowerCAmelCase:Union[str, Any]									=  []
											for i, rect in enumerate(a__):
																rect.set_stroke(a__)
																# target = fill.copy().set_fill(YELLOW, opacity=0.7)
																# target.move_to(rect)
																# self.add(target)

																_lowerCAmelCase:Dict									=  Rectangle(height=0.46 / 4     ,width=0.46 / 3).set_stroke(width=0.0).set_fill(a__     ,opacity=0.7)

																if i == 0:
																					cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT)     ,buff=0.02     ,direction=a__)
																					cpu_target.set_x(cpu_target.get_x() + 0.1)
																elif i == 3:
																					cpu_target.next_to(cpu_targs[0]     ,direction=a__     ,buff=0.0)
																else:
																					cpu_target.next_to(cpu_targs[i - 1]     ,direction=a__     ,buff=0.0)
																self.add(a__)
																cpu_targs.append(a__)

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

											_lowerCAmelCase:Union[str, Any]									=  Text('''Loaded Checkpoint'''     ,font_size=24)
											_lowerCAmelCase:Optional[Any]									=  Group(a__     ,a__).arrange(a__     ,aligned_edge=a__     ,buff=0.4)
											checkpoint.move_to([3, 0.5, 0])

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

											_lowerCAmelCase:Any									=  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(a__     ,a__)

											_lowerCAmelCase:Tuple									=  MarkupText(
											    F'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'     ,font_size=18     ,)

											blue_text.next_to(a__     ,DOWN * 2.4     ,aligned_edge=key_text.get_left())

											_lowerCAmelCase:List[str]									=  MarkupText(
											    F'Next, a <i><span fgcolor="{BLUE}">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor="{BLUE}">single shard</span>.'     ,font_size=24     ,)
											step_a.move_to([2, 2, 0])
											self.play(Write(a__)     ,Write(a__))

											self.play(Write(a__     ,run_time=1)     ,Create(a__     ,run_time=1))

											_lowerCAmelCase:str									=  []
											_lowerCAmelCase:int									=  []
											for i, rect in enumerate(a__):
																_lowerCAmelCase:Tuple									=  fill.copy().set_fill(a__     ,opacity=0.7)
																target.move_to(a__)
																first_animations.append(GrowFromCenter(a__     ,run_time=1))

																_lowerCAmelCase:List[str]									=  target.copy()
																cpu_target.generate_target()
																if i < 5:
																					cpu_target.target.move_to(cpu_left_col_base[i + 1])
																else:
																					cpu_target.target.move_to(cpu_right_col_base[i - 5])
																second_animations.append(MoveToTarget(a__     ,run_time=1.5))

											self.play(*a__)
											self.play(*a__)
											self.wait()












"""simple docstring"""






import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp





def 							UpperCAmelCase		(		):
					_lowerCAmelCase:Optional[int]									=  ArgumentParser(
					    description=(
					        '''PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes'''
					    )    )

					# Optional arguments for the launch helper
					parser.add_argument('''--num_cores'''      ,       type=snake_case      ,       default=1      ,       help='''Number of TPU cores to use (1 or 8).'''    )

					# positional
					parser.add_argument(
					    '''training_script'''      ,       type=snake_case      ,       help=(
					        '''The full path to the single TPU training '''
					        '''program/script to be launched in parallel, '''
					        '''followed by all the arguments for the '''
					        '''training script'''
					    )      ,       )

					# rest from the training program
					parser.add_argument('''training_script_args'''      ,       nargs=snake_case    )

					return parser.parse_args()





def 							UpperCAmelCase		(		):
					_lowerCAmelCase:int									=  parse_args()

					# Import training_script as a module.
					_lowerCAmelCase:Optional[int]									=  Path(args.training_script    )
					sys.path.append(str(script_fpath.parent.resolve()    )    )
					_lowerCAmelCase:Optional[int]									=  script_fpath.stem
					_lowerCAmelCase:Any									=  importlib.import_module(snake_case    )

					# Patch sys.argv
					_lowerCAmelCase:Dict									=  [args.training_script] + args.training_script_args + ['''--tpu_num_cores''', str(args.num_cores    )]

					xmp.spawn(mod._mp_fn      ,       args=()      ,       nprocs=args.num_cores    )


if __name__ == "__main__":
				main()






