code: 
  
  
  
  
  
  import json
  import os
  from collections.abc import Mapping
  from typing import Dict, List, Optional, Union
  
  import numpy as np
  import torch
  
  from .imports import is_safetensors_available
  def   _A (	SCREAMING_SNAKE_CASE  :     Optional[Any] , SCREAMING_SNAKE_CASE  :     Union[str, Any] , SCREAMING_SNAKE_CASE  :     List[str] , SCREAMING_SNAKE_CASE  :     List[str]=None		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      List[Any]            =None
    # Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
    if str(weight.dtype		) == "torch.bfloat16":
      # Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
      a__			:      List[Any]            =weight.view(torch.intaa		)
      a__			:      Optional[Any]            ="bfloat16"
    a__			:      str            =weight.cpu().numpy()
    a__			:      str            =os.path.join(SCREAMING_SNAKE_CASE , f'''{weight_name}.dat'''		)
    if index is not None:
      if dtype is None:
        a__			:      Any            =str(array.dtype		)
      a__			:      Union[str, Any]            ={"dtype": dtype, "shape": list(array.shape		)}
    if array.ndim == 0:
      a__			:      int            =array[None]
    a__			:      List[Any]            =np.memmap(SCREAMING_SNAKE_CASE , dtype=array.dtype , mode="w+" , shape=array.shape		)
    a__			:      Tuple            =array[:]
    file_array.flush()
    return index
  def   _A (	SCREAMING_SNAKE_CASE  :     List[str] , SCREAMING_SNAKE_CASE  :     List[str]		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      Dict            =tuple(weight_info["shape"]		)
    if shape == ():
      # NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
      a__			:      str            =(1,)
  
    a__			:      Union[str, Any]            =weight_info["dtype"]
    if dtype == "bfloat16":
      # NumPy does not support bfloat16 so this was saved as a int16
      a__			:      Any            ="int16"
  
    a__			:      Dict            =np.memmap(SCREAMING_SNAKE_CASE , dtype=SCREAMING_SNAKE_CASE , shape=SCREAMING_SNAKE_CASE , mode="r"		)
  
    if len(weight_info["shape"]		) == 0:
      a__			:      List[Any]            =weight[0]
    a__			:      Tuple            =torch.tensor(SCREAMING_SNAKE_CASE		)
    if weight_info["dtype"] == "bfloat16":
      a__			:      Dict            =weight.view(torch.bfloataa		)
  
    return weight
  def   _A (	SCREAMING_SNAKE_CASE  :     List[Any] , SCREAMING_SNAKE_CASE  :     Union[str, Any]		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    if index is None or len(SCREAMING_SNAKE_CASE		) == 0:
      # Nothing to save
      return
  
    a__			:      Optional[int]            =os.path.join(SCREAMING_SNAKE_CASE , "index.json"		)
    if os.path.isfile(SCREAMING_SNAKE_CASE		):
      with open(SCREAMING_SNAKE_CASE , "r" , encoding="utf-8"		) as f:
        a__			:      Any            =json.load(SCREAMING_SNAKE_CASE		)
    else:
      a__			:      Any            ={}
    current_index.update(SCREAMING_SNAKE_CASE		)
  
    with open(SCREAMING_SNAKE_CASE , "w" , encoding="utf-8"		) as f:
      json.dump(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , indent=2		)
  def   _A (	SCREAMING_SNAKE_CASE  :     Union[str, os.PathLike] , SCREAMING_SNAKE_CASE  :     Dict[str, torch.Tensor]		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    os.makedirs(SCREAMING_SNAKE_CASE , exist_ok=SCREAMING_SNAKE_CASE		)
    a__			:      Any            ={}
    for name, parameter in state_dict.items():
      a__			:      List[str]            =offload_weight(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , index=SCREAMING_SNAKE_CASE		)
  
    # Update index
    save_offload_index(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE		)
  
  
  
  
  
  
  
  class        __lowerCAmelCase (					UpperCamelCase__):
  
  
  
  
  
  
  
        def __init__(					self	,   lowerCAmelCase__	,   lowerCAmelCase__							)     ->		Optional[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Optional[Any]            =dataset
          a__			:      str            =prefix
  
  
  
  
  
  
  
        def __getitem__(					self	,   lowerCAmelCase__							)     ->		str:
  
  
  
          '''simple docstring'''
  
  
  
  
          return self.dataset[F'''{self.prefix}{key}''']
  
  
  
  
  
  
  
        def __iter__(					self							)     ->		Union[str, Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          return iter([key for key in self.dataset if key.startswith(self.prefix							)]							)
  
  
  
  
  
  
  
        def __len__(					self							)     ->		Optional[int]:
  
  
  
          '''simple docstring'''
  
  
  
  
          return len(self.dataset							)
  
  
  
  
  
  
  
  class        __lowerCAmelCase (					UpperCamelCase__):
  
  
  
  
  
  
  
        def __init__(					self	,   lowerCAmelCase__ = None	,   lowerCAmelCase__ = None	,   lowerCAmelCase__ = None	,   lowerCAmelCase__=None	,   )     ->		Any:
  
  
  
          '''simple docstring'''
  
  
  
  
          if state_dict is None and save_folder is None:
            raise ValueError("Need either a `state_dict` or a `save_folder` containing offloaded weights."							)
  
          a__			:      int            ={} if state_dict is None else state_dict
          a__			:      int            =save_folder
          if index is None and save_folder is not None:
            with open(os.path.join(lowerCAmelCase__	,   "index.json"							)							) as f:
              a__			:      str            =json.load(lowerCAmelCase__							)
          a__			:      Tuple            ={} if index is None else index
          a__			:      Union[str, Any]            =list(self.state_dict.keys()							)
          self.all_keys.extend([key for key in self.index if key not in self.all_keys]							)
          a__			:      List[Any]            =device
  
  
  
  
  
  
  
        def __getitem__(					self	,   lowerCAmelCase__							)     ->		str:
  
  
  
          '''simple docstring'''
  
  
  
  
          if key in self.state_dict:
            return self.state_dict[key]
          a__			:      str            =self.index[key]
          if weight_info.get("safetensors_file"							) is not None:
            if not is_safetensors_available():
              raise ImportError("These offloaded weights require the use of safetensors: `pip install safetensors`."							)
  
            from safetensors import safe_open
  
            a__			:      Optional[Any]            ="cpu" if self.device is None else self.device
            with safe_open(weight_info["safetensors_file"]	,   framework="pt"	,   device=lowerCAmelCase__							) as f:
              a__			:      Optional[int]            =f.get_tensor(weight_info.get("weight_name"	,   lowerCAmelCase__							)							)
  
            if "dtype" in weight_info:
              return tensor.to(getattr(lowerCAmelCase__	,   weight_info["dtype"]							)							)
            else:
              return tensor
  
          a__			:      Any            =os.path.join(self.save_folder	,   F'''{key}.dat'''							)
          return load_offloaded_weight(lowerCAmelCase__	,   lowerCAmelCase__							)
  
  
  
  
  
  
  
        def __iter__(					self							)     ->		Optional[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          return iter(self.all_keys							)
  
  
  
  
  
  
  
        def __len__(					self							)     ->		Optional[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          return len(self.all_keys							)
  def   _A (	SCREAMING_SNAKE_CASE  :     Dict[str, torch.Tensor] , SCREAMING_SNAKE_CASE  :     List[str]		):
  
  
  
  
  
    """simple docstring"""
  
  
  
    a__			:      Optional[Any]            ={}
    for module_name in submodule_names:
      # We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
      # submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
      result.update(
          {
              key: param
              for key, param in state_dict.items()
              if key == module_name or key.startswith(module_name + "."		)
          }		)
    return result
  
  
  
  
  
  
  
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()
  
  
  
  
  
  
  