content: import gc
  import os
  import re
  import time
  from collections import defaultdict
  from contextlib import contextmanager
  
  import psutil
  import torch
  from packaging import version
  
  from .config import config_to_primitive
  from .core import debug, find, info, warn
  from .typing import *
  
  
  def parse_version(ver: str):
      return version.parse(ver)
  
  
  def get_rank():
      # SLURM_PROCID can be set even if SLURM is not managing the multiprocessing,
      # therefore LOCAL_RANK needs to be checked first
      rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK")
      for key in rank_keys:
          rank = os.environ.get(key)
          if rank is not None:
              return int(rank)
      return 0
  
  
  def get_device():
      return torch.device(f"cuda:{get_rank()}")
  
  
  def load_module_weights(
      path, module_name=None, ignore_modules=None, mapping=None, map_location=None
  ) -> Tuple[dict, int, int]:
      if module_name is not None and ignore_modules is not None:
          raise ValueError("module_name and ignore_modules cannot be both set")
      if map_location is None:
          map_location = get_device()
  
      ckpt = torch.load(path, map_location=map_location)
      state_dict = ckpt["state_dict"]
  
      if mapping is not None:
          state_dict_to_load = {}
          for k, v in state_dict.items():
              if any([k.startswith(m["to"]) for m in mapping]):
                  pass
              else:
                  state_dict_to_load[k] = v
          for k, v in state_dict.items():
              for m in mapping:
                  if k.startswith(m["from"]):
                      k_dest = k.replace(m["from"], m["to"])
                      info(f"Mapping {k} => {k_dest}")
                      state_dict_to_load[k_dest] = v.clone()
          state_dict = state_dict_to_load
  
      state_dict_to_load = state_dict
  
      if ignore_modules is not None:
          state_dict_to_load = {}
          for k, v in state_dict.items():
              ignore = any(
                  [k.startswith(ignore_module + ".") for ignore_module in ignore_modules]
              )
              if ignore:
                  continue
              state_dict_to_load[k] = v
  
      if module_name is not None:
          state_dict_to_load = {}
          for k, v in state_dict.items():
              m = re.match(rf"^{module_name}\.(.*)$", k)
              if m is None:
                  continue
              state_dict_to_load[m.group(1)] = v
  
      return state_dict_to_load, ckpt["epoch"], ckpt["global_step"]
  
  
  def C(value: Any, epoch: int, global_step: int) -> float:
      if isinstance(value, int) or isinstance(value, float):
          pass
      else:
          value = config_to_primitive(value)
          if not isinstance(value, list):
              raise TypeError("Scalar specification only supports list, got", type(value))
          if len(value) == 3:
              value = [0] + value
          assert len(value) == 4
          start_step, start_value, end_value, end_step = value
          if isinstance(end_step, int):
              current_step = global_step
              value = start_value + (end_value - start_value) * max(
                  min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0
              )
          elif isinstance(end_step, float):
              current_step = epoch
              value = start_value + (end_value - start_value) * max(
                  min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0
              )
      return value
  
  
  def cleanup():
      gc.collect()
      torch.cuda.empty_cache()
      try:
          import tinycudann as tcnn
  
          tcnn.free_temporary_memory()
      except:
          pass
  
  
  def finish_with_cleanup(func: Callable):
      def wrapper(*args, **kwargs):
          out = func(*args, **kwargs)
          cleanup()
          return out
  
      return wrapper
  
  
  def _distributed_available():
      return torch.distributed.is_available() and torch.distributed.is_initialized()
  
  
  def barrier():
      if not _distributed_available():
          return
      else:
          torch.distributed.barrier()
  
  
  def broadcast(tensor, src=0):
      if not _distributed_available():
          return tensor
      else:
          torch.distributed.broadcast(tensor, src=src)
          return tensor
  
  
  def enable_gradient(model, enabled: bool = True) -> None:
      for param in model.parameters():
          param.requires_grad_(enabled)
  
  
  class TimeRecorder:
      _instance = None
  
      def __init__(self):
          self.items = {}
          self.accumulations = defaultdict(list)
          self.time_scale = 1000.0  # ms
          self.time_unit = "ms"
          self.enabled = False
  
      def __new__(cls):
          # singleton
          if cls._instance is None:
              cls._instance = super(TimeRecorder, cls).__new__(cls)
          return cls._instance
  
      def enable(self, enabled: bool) -> None:
          self.enabled = enabled
  
      def start(self, name: str) -> None:
          if not self.enabled:
              return
          torch.cuda.synchronize()
          self.items[name] = time.time()
  
      def end(self, name: str, accumulate: bool = False) -> float:
          if not self.enabled or name not in self.items:
              return
          torch.cuda.synchronize()
          start_time = self.items.pop(name)
          delta = time.time() - start_time
          if accumulate:
              self.accumulations[name].append(delta)
          t = delta * self.time_scale
          info(f"{name}: {t:.2f}{self.time_unit}")
  
      def get_accumulation(self, name: str, average: bool = False) -> float:
          if not self.enabled or name not in self.accumulations:
              return
          acc = self.accumulations.pop(name)
          total = sum(acc)
          if average:
              t = total / len(acc) * self.time_scale
          else:
              t = total * self.time_scale
          info(f"{name} for {len(acc)} times: {t:.2f}{self.time_unit}")
  
  
  ### global time recorder
  time_recorder = TimeRecorder()
  
  
  @contextmanager
  def time_recorder_enabled():
      enabled = time_recorder.enabled
      time_recorder.enable(enabled=True)
      try:
          yield
      finally:
          time_recorder.enable(enabled=enabled)
  
  
  def show_vram_usage(name):
      available, total = torch.cuda.mem_get_info()
      used = total - available
      print(
          f"{name}: {used / 1024**2:.1f}MB, {psutil.Process(os.getpid()).memory_info().rss / 1024**2:.1f}MB"
      )
  