




import os
import sys
import tempfile

import torch

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





def __UpperCamelCase			(       _A ,  _A=() ,  _A=None ,  _A="no" ,  _A="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''' ,  _A  ) 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(_A ,  distributed_type='''TPU'''  )
												print(f"Launching a training on {num_processes} TPU cores."  )
												xmp.spawn(_A ,  args=_A ,  nprocs=_A ,  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(*_A  )
						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=_A ,  master_addr='''127.0.01''' ,  master_port=_A ,  mixed_precision=_A  ):
																								lowerCAmelCase_      			=  PrepareForLaunch(_A ,  distributed_type='''MULTI_GPU'''  )
																								print(f"Launching training on {num_processes} GPUs."  )
																								try:
																														start_processes(_A ,  args=_A ,  nprocs=_A ,  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(*_A  )






def __UpperCamelCase			(       _A ,  _A=() ,  _A=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=_A ,  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(_A ,  debug=_A  )
																		start_processes(_A ,  args=_A ,  nprocs=_A ,  start_method='''fork'''  )










from functools import lru_cache





@lru_cache
def __UpperCamelCase			(       _A  ):
						if num < 0:
												raise ValueError('''Number should not be negative.'''  )

						return 1 if num in (0, 1) else num * factorial(num - 1  )


if __name__ == "__main__":
						import doctest

						doctest.testmod()




