





import os
import sys
import tempfile

import torch

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




def _a (  SCREAMING_SNAKE_CASE__      :  Tuple			,				SCREAMING_SNAKE_CASE__      :  Any=()			,				SCREAMING_SNAKE_CASE__      :  Optional[Any]=None			,				SCREAMING_SNAKE_CASE__      :  Tuple="no"			,				SCREAMING_SNAKE_CASE__      :  Optional[int]="29500"       )	->    Optional[Any]:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      Any     =						False
							SCREAMING_SNAKE_CASE__       :      int     =						False
							if any(key.startswith("KAGGLE"       ) for key in os.environ.keys()       ):
														SCREAMING_SNAKE_CASE__       :      Tuple     =						True
							elif "IPython" in sys.modules:
														SCREAMING_SNAKE_CASE__       :      Dict     =						"google.colab" in str(sys.modules["IPython"].get_ipython()       )

							try:
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						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"			,				__SCREAMING_SNAKE_CASE       ) 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:
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						8

														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						PrepareForLaunch(__SCREAMING_SNAKE_CASE			,				distributed_type="TPU"       )
														print(f'''Launching a training on {num_processes} TPU cores.'''       )
														xmp.spawn(__SCREAMING_SNAKE_CASE			,				args=__SCREAMING_SNAKE_CASE			,				nprocs=__SCREAMING_SNAKE_CASE			,				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(*__SCREAMING_SNAKE_CASE       )
							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=__SCREAMING_SNAKE_CASE			,				master_addr="127.0.01"			,				master_port=__SCREAMING_SNAKE_CASE			,				mixed_precision=__SCREAMING_SNAKE_CASE       ):
																												SCREAMING_SNAKE_CASE__       :      Optional[int]     =						PrepareForLaunch(__SCREAMING_SNAKE_CASE			,				distributed_type="MULTI_GPU"       )
																												print(f'''Launching training on {num_processes} GPUs.'''       )
																												try:
																																			start_processes(__SCREAMING_SNAKE_CASE			,				args=__SCREAMING_SNAKE_CASE			,				nprocs=__SCREAMING_SNAKE_CASE			,				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():
																												SCREAMING_SNAKE_CASE__       :      List[Any]     =						"1"
																												print("Launching training on MPS."       )
																					elif torch.cuda.is_available():
																												print("Launching training on one GPU."       )
																					else:
																												print("Launching training on CPU."       )
																					function(*__SCREAMING_SNAKE_CASE       )







def _a (  SCREAMING_SNAKE_CASE__      :  int			,				SCREAMING_SNAKE_CASE__      :  Optional[int]=()			,				SCREAMING_SNAKE_CASE__      :  str=2       )	->    List[Any]:

							'''simple docstring'''






							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=__SCREAMING_SNAKE_CASE			,				master_addr="127.0.01"			,				master_port="29500"			,				accelerate_mixed_precision="no"			,				accelerate_debug_rdv_file=tmp_file.name			,				accelerate_use_cpu="yes"			,				):
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						PrepareForLaunch(__SCREAMING_SNAKE_CASE			,				debug=__SCREAMING_SNAKE_CASE       )
																					start_processes(__SCREAMING_SNAKE_CASE			,				args=__SCREAMING_SNAKE_CASE			,				nprocs=__SCREAMING_SNAKE_CASE			,				start_method="fork"       )













def _a (  SCREAMING_SNAKE_CASE__      :  int = 50_00_00_00       )	->    int:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      str     =						set()
							SCREAMING_SNAKE_CASE__       :      Dict     =						int((limit - 24) ** (1 / 2)       )

							SCREAMING_SNAKE_CASE__       :      Any     =						set(range(3			,				prime_square_limit + 1			,				2       )       )
							primes.add(2       )
							for p in range(3			,				prime_square_limit + 1			,				2       ):
														if p not in primes:
																					continue
														primes.difference_update(set(range(p * p			,				prime_square_limit + 1			,				SCREAMING_SNAKE_CASE__       )       )       )

							for primea in primes:
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						primea * primea
														for primea in primes:
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						primea * primea * primea
																					if square + cube >= limit - 16:
																												break
																					for primea in primes:
																												SCREAMING_SNAKE_CASE__       :      List[str]     =						primea * primea * primea * primea
																												SCREAMING_SNAKE_CASE__       :      Optional[int]     =						square + cube + tetr
																												if total >= limit:
																																			break
																												ret.add(SCREAMING_SNAKE_CASE__       )

							return len(SCREAMING_SNAKE_CASE__       )


if __name__ == "__main__":
		print(f"{solution() = }")






