code: 
  
  
  
  
  
  from typing import List, Optional, Tuple, Union
  
  from . import logging
  from .import_utils import is_torch_available, is_torch_version
  
  
  if is_torch_available():
         import torch
  
  UpperCAmelCase			:			str	  = logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  try:
         from torch._dynamo import allow_in_graph as maybe_allow_in_graph
  except (ImportError, ModuleNotFoundError):
         def   _A (	cls  :     Any		):
  
  
  
  
  
           """simple docstring"""
  
  
  
           return cls
  def   _A (	SCREAMING_SNAKE_CASE  :     Union[Tuple, List] , SCREAMING_SNAKE_CASE  :     Optional[Union[List["torch.Generator"], "torch.Generator"]] = None , SCREAMING_SNAKE_CASE  :     Optional["torch.device"] = None , SCREAMING_SNAKE_CASE  :     Optional["torch.dtype"] = None , SCREAMING_SNAKE_CASE  :     Optional["torch.layout"] = None , ):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      Dict            =device
    a__			:      Tuple            =shape[0]
  
    a__			:      Union[str, Any]            =layout or torch.strided
    a__			:      Dict            =device or torch.device("cpu"		)
  
    if generator is not None:
      a__			:      Optional[int]            =generator.device.type if not isinstance(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		) else generator[0].device.type
      if gen_device_type != device.type and gen_device_type == "cpu":
        a__			:      List[str]            ="cpu"
        if device != "mps":
          logger.info(
              f'''The passed generator was created on \'cpu\' even though a tensor on {device} was expected.'''
              f''' Tensors will be created on \'cpu\' and then moved to {device}. Note that one can probably'''
              f''' slighly speed up this function by passing a generator that was created on the {device} device.'''		)
      elif gen_device_type != device.type and gen_device_type == "cuda":
        raise ValueError(f'''Cannot generate a {device} tensor from a generator of type {gen_device_type}.'''		)
  
    if isinstance(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		):
      a__			:      Tuple            =(1,) + shape[1:]
      a__			:      str            =[
          torch.randn(SCREAMING_SNAKE_CASE , generator=generator[i] , device=SCREAMING_SNAKE_CASE , dtype=SCREAMING_SNAKE_CASE , layout=SCREAMING_SNAKE_CASE		)
          for i in range(SCREAMING_SNAKE_CASE		)
      ]
      a__			:      Dict            =torch.cat(SCREAMING_SNAKE_CASE , dim=0		).to(SCREAMING_SNAKE_CASE		)
    else:
      a__			:      Optional[Any]            =torch.randn(SCREAMING_SNAKE_CASE , generator=SCREAMING_SNAKE_CASE , device=SCREAMING_SNAKE_CASE , dtype=SCREAMING_SNAKE_CASE , layout=SCREAMING_SNAKE_CASE		).to(SCREAMING_SNAKE_CASE		)
  
    return latents
  
  
  
  
  
  def   _A (	SCREAMING_SNAKE_CASE  :     str		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    if is_torch_version("<" , "2.0.0"		) or not hasattr(SCREAMING_SNAKE_CASE , "_dynamo"		):
      return False
    return isinstance(SCREAMING_SNAKE_CASE , torch._dynamo.eval_frame.OptimizedModule		)
  
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  import importlib
  import sys
  from argparse import REMAINDER, ArgumentParser
  from pathlib import Path
  
  import torch_xla.distributed.xla_multiprocessing as xmp
  def   _A (	):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      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=SCREAMING_SNAKE_CASE , default=1 , help="Number of TPU cores to use (1 or 8)."		)
  
    # positional
    parser.add_argument(
        "training_script" , type=SCREAMING_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=SCREAMING_SNAKE_CASE		)
  
    return parser.parse_args()
  def   _A (	):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      Optional[Any]            =parse_args()
  
    # Import training_script as a module.
    a__			:      Tuple            =Path(args.training_script		)
    sys.path.append(str(script_fpath.parent.resolve()		)		)
    a__			:      Tuple            =script_fpath.stem
    a__			:      Any            =importlib.import_module(SCREAMING_SNAKE_CASE		)
  
    # Patch sys.argv
    a__			:      List[str]            =[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()
  
  
  
  
  
  
  