
"""simple docstring"""


import os
import sys
import tempfile

import torch

from .state import AcceleratorState
from .utils import PrecisionType, PrepareForLaunch, is_mps_available, patch_environment






def  lowercase						(	lowerCAmelCase__	,lowerCAmelCase__=()	,lowerCAmelCase__=None	,lowerCAmelCase__="no"	,lowerCAmelCase__="29500"						):
			lowerCamelCase_					=  False
			lowerCamelCase_					=  False
			if any(key.startswith('''KAGGLE'''						) for key in os.environ.keys()						):
						lowerCamelCase_					=  True
			elif "IPython" in sys.modules:
						lowerCamelCase_					=  '''google.colab''' in str(sys.modules['''IPython'''].get_ipython()						)

			try:
						lowerCamelCase_					=  PrecisionType(mixed_precision.lower()						)
			except ValueError:
						raise ValueError(
						    f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}."						)

			if (in_colab or in_kaggle) and (os.environ.get('''TPU_NAME'''	,lowerCAmelCase__						) is not None):
						# TPU launch
						import torch_xla.distributed.xla_multiprocessing as xmp

						if len(AcceleratorState._shared_state						) > 0:
									raise ValueError(
									    '''To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside '''
									    '''your training function. Restart your notebook and make sure no cells initializes an '''
									    '''`Accelerator`.'''						)
						if num_processes is None:
									lowerCamelCase_					=  8

						lowerCamelCase_					=  PrepareForLaunch(lowerCAmelCase__	,distributed_type='''TPU'''						)
						print(f"Launching a training on {num_processes} TPU cores."						)
						xmp.spawn(lowerCAmelCase__	,args=lowerCAmelCase__	,nprocs=lowerCAmelCase__	,start_method='''fork'''						)
			elif in_colab:
						# No need for a distributed launch otherwise as it's either CPU or one GPU.
						if torch.cuda.is_available():
									print('''Launching training on one GPU.'''						)
						else:
									print('''Launching training on one CPU.'''						)
						function(*lowerCAmelCase__						)
			else:
						if num_processes is None:
									raise ValueError(
									    '''You have to specify the number of GPUs you would like to use, add `num_processes=...` to your call.'''						)

						if num_processes > 1:
									# Multi-GPU launch
									from torch.multiprocessing import start_processes
									from torch.multiprocessing.spawn import ProcessRaisedException

									if len(AcceleratorState._shared_state						) > 0:
												raise ValueError(
												    '''To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized '''
												    '''inside your training function. Restart your notebook and make sure no cells initializes an '''
												    '''`Accelerator`.'''						)

									if torch.cuda.is_initialized():
												raise ValueError(
												    '''To launch a multi-GPU training from your notebook, you need to avoid running any instruction '''
												    '''using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA '''
												    '''function.'''						)

									# torch.distributed will expect a few environment variable to be here. We set the ones common to each
									# process here (the other ones will be set be the launcher).
									with patch_environment(
									    world_size=lowerCAmelCase__	,master_addr='''127.0.01'''	,master_port=lowerCAmelCase__	,mixed_precision=lowerCAmelCase__						):
												lowerCamelCase_					=  PrepareForLaunch(lowerCAmelCase__	,distributed_type='''MULTI_GPU'''						)
												print(f"Launching training on {num_processes} GPUs."						)
												try:
															start_processes(lowerCAmelCase__	,args=lowerCAmelCase__	,nprocs=lowerCAmelCase__	,start_method='''fork'''						)
												except ProcessRaisedException as e:
															if "Cannot re-initialize CUDA in forked subprocess" in e.args[0]:
																		raise RuntimeError(
																		    '''CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. '''
																		    '''This likely stems from an outside import causing issues once the `notebook_launcher()` is called. '''
																		    '''Please review your imports and test them when running the `notebook_launcher()` to identify '''
																		    '''which one is problematic.'''						) from e

						else:
									# No need for a distributed launch otherwise as it's either CPU, GPU or MPS.
									if is_mps_available():
												lowerCamelCase_					=  '''1'''
												print('''Launching training on MPS.'''						)
									elif torch.cuda.is_available():
												print('''Launching training on one GPU.'''						)
									else:
												print('''Launching training on CPU.'''						)
									function(*lowerCAmelCase__						)






def  lowercase						(	lowerCAmelCase__	,lowerCAmelCase__=()	,lowerCAmelCase__=2						):
			from torch.multiprocessing import start_processes

			with tempfile.NamedTemporaryFile() as tmp_file:
						# torch.distributed will expect a few environment variable to be here. We set the ones common to each
						# process here (the other ones will be set be the launcher).
						with patch_environment(
						    world_size=lowerCAmelCase__	,master_addr='''127.0.01'''	,master_port='''29500'''	,accelerate_mixed_precision='''no'''	,accelerate_debug_rdv_file=tmp_file.name	,accelerate_use_cpu='''yes'''	,):
									lowerCamelCase_					=  PrepareForLaunch(lowerCAmelCase__	,debug=lowerCAmelCase__						)
									start_processes(lowerCAmelCase__	,args=lowerCAmelCase__	,nprocs=lowerCAmelCase__	,start_method='''fork'''						)










from manim import *

class 							__SCREAMING_SNAKE_CASE       (  _a							):






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      Rectangle(height=0.5	,			width=0.5	)
														UpperCamelCase__					  =      Rectangle(height=0.25	,			width=0.25	)
														UpperCamelCase__					  =      Rectangle(height=0.46	,			width=0.46	).set_stroke(width=0	)

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

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

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

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

														UpperCamelCase__					  =      []
														UpperCamelCase__					  =      []
														UpperCamelCase__					  =      []

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

																					UpperCamelCase__					  =      Rectangle(height=0.46 / 4	,			width=0.46 / 3	).set_stroke(width=0.0	).set_fill(__lowerCAmelCase	,			opacity=0.7	)

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

														self.add(*__lowerCAmelCase	,			*__lowerCAmelCase	,			*__lowerCAmelCase	)

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

														UpperCamelCase__					  =      Text("""Loaded Checkpoint"""	,			font_size=24	)
														UpperCamelCase__					  =      Group(__lowerCAmelCase	,			__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0.5	,			aligned_edge=__lowerCAmelCase	)
														checkpoint.move_to([3, 0.5, 0]	)

														self.add(__lowerCAmelCase	)

														UpperCamelCase__					  =      []
														UpperCamelCase__					  =      []

														for i, rect in enumerate(__lowerCAmelCase	):
																					UpperCamelCase__					  =      fill.copy().set_fill(__lowerCAmelCase	,			opacity=0.7	)
																					target.move_to(__lowerCAmelCase	)
																					ckpt_arr.append(__lowerCAmelCase	)

																					UpperCamelCase__					  =      target.copy()
																					if i < 5:
																												cpu_target.move_to(cpu_left_col_base[i + 1]	)
																					else:
																												cpu_target.move_to(cpu_right_col_base[i - 5]	)
																					ckpt_cpu_arr.append(__lowerCAmelCase	)
														self.add(*__lowerCAmelCase	,			*__lowerCAmelCase	)

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

														UpperCamelCase__					  =      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	)

														UpperCamelCase__					  =      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	)

														UpperCamelCase__					  =      MarkupText(
														    f"""Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device."""	,			font_size=24	,			)
														step_a.move_to([2, 2, 0]	)

														UpperCamelCase__					  =      [meta_mem.copy() for i in range(6	)]
														UpperCamelCase__					  =      [meta_mem.copy() for i in range(6	)]
														UpperCamelCase__					  =      VGroup(*__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)
														UpperCamelCase__					  =      VGroup(*__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)
														UpperCamelCase__					  =      VGroup(__lowerCAmelCase	,			__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)
														UpperCamelCase__					  =      Text("""Disk"""	,			font_size=24	)
														UpperCamelCase__					  =      Group(__lowerCAmelCase	,			__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0.5	,			aligned_edge=__lowerCAmelCase	)
														disk.move_to([-4.0, -1.25, 0]	)
														self.play(Write(__lowerCAmelCase	,			run_time=3	)	,			Write(__lowerCAmelCase	,			run_time=1	)	,			Create(__lowerCAmelCase	,			run_time=1	)	)

														UpperCamelCase__					  =      []
														for i, rect in enumerate(__lowerCAmelCase	):
																					UpperCamelCase__					  =      rect.copy()
																					target.generate_target()
																					target.target.move_to(disk_left_col_base[i]	).scale(0.5	)
																					animations.append(MoveToTarget(__lowerCAmelCase	,			run_time=1.5	)	)
														self.play(*__lowerCAmelCase	)

														self.play(FadeOut(__lowerCAmelCase	)	)

														UpperCamelCase__					  =      MarkupText(f"""Then, the checkpoint is removed from memory\nthrough garbage collection."""	,			font_size=24	)
														step_a.move_to([2, 2, 0]	)

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

														self.play(
														    FadeOut(__lowerCAmelCase	,			__lowerCAmelCase	,			*__lowerCAmelCase	,			*__lowerCAmelCase	)	,			)

														self.wait()




