code: def clean_device_map(
      device_map: Dict[str, Union[int, str, torch.device]],
      module_name: str = "",
  ):
      """
      Cleans a device_map by grouping all submodules that go on the same device together.
      """
      # Get the value of the current module and if there is only one split across several keys, regroup it.
      prefix = "" if module_name == "" else f"{module_name}."
      values = [
          v for k, v in device_map.items() if k.startswith(prefix)
      ]
      if len(set(values)) == 1 and len(values) > 1:
          for k in [k for k in device_map if k.startswith(prefix)]:
              del device_map[k]
          device_map[module_name] = values[0]
  
      # Recurse over the children
      children_modules = [
          k
          for k in device_map.keys()
          if k.startswith(prefix) and len(k) > len(module_name)
      ]
      idx = (
          len(module_name.split(".")) + 1 if len(module_name) > 0 else 1
      )
      children_modules = set(
          ".".join(k.split(".")[:idx]) for k in children_modules
      )
      for child in children_modules:
          clean_device_map(device_map, module_name=child)
  
      return device_map
  
imports: import contextlib
  import gc
  import inspect
  import json
  import logging
  import os
  import re
  import shutil
  import tempfile
  from collections import OrderedDict, defaultdict
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  import torch.nn
  from ..state import AcceleratorState
  from .constants import SAFE_WEIGHTS_NAME, WEIGHTS_NAME
  from .dataclasses import AutocastKwargs, CustomDtype, DistributedType
  from .imports import is_mps_available, is_npu_available, is_xpu_available
  from .offload import load_offloaded_weight, offload_weight, save_offload_index
  from .tqdm import is_tqdm_available, tqdm
  from safetensors import safe_open
  from safetensors.torch import load_file
label_desc: Use image node_id 18 for calling a global function with example usage: clean_device_map(device_map, module_name) and returns: device_map