


from __future__ import annotations






def    _UpperCAmelCase	(   __lowerCamelCase       )		->list[int]:


 '''simple docstring'''
 assert (
     isinstance(__lowerCamelCase ,     __lowerCamelCase       ) and mask > 0
 ), f'''mask needs to be positive integer, your input {mask}'''
 snake_case   : Tuple	  =      []
 snake_case   : str	  =      mask

 while submask:
  all_submasks.append(__lowerCamelCase       )
  snake_case   : Optional[Any]	  =      (submask - 1) & mask

 return all_submasks


if __name__ == "__main__":
  import doctest

  doctest.testmod()










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

import torch_xla.distributed.xla_multiprocessing as xmp






def    _UpperCAmelCase	(   )		->str:


 '''simple docstring'''
 snake_case   : List[str]	  =      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=__lowerCamelCase ,     default=1 ,     help='''Number of TPU cores to use (1 or 8).'''       )

 # positional
 parser.add_argument(
     '''training_script''' ,     type=__lowerCamelCase ,     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=__lowerCamelCase       )

 return parser.parse_args()






def    _UpperCAmelCase	(   )		->Optional[int]:


 '''simple docstring'''
 snake_case   : Optional[Any]	  =      parse_args()

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

 # Patch sys.argv
 snake_case   : int	  =      [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()






