code: def attach_align_device_hook_on_blocks(
      module: nn.Module,
      execution_device: Optional[
          Union[torch.device, Dict[str, torch.device]]
      ] = None,
      offload: Union[bool, Dict[str, bool]] = False,
      weights_map: Mapping = None,
      offload_buffers: bool = False,
      module_name: str = "",
      skip_keys: Optional[Union[str, List[str]]] = None,
      preload_module_classes: Optional[List[str]] = None,
  ):
      """
      Attaches `AlignDevicesHook` to all blocks of a given model as needed.
  
      Args:
          module (`torch.nn.Module`):
              The module where we want to attach the hooks.
          execution_device (`torch.device` or `Dict[str, torch.device]`, *optional*):
              The device on which inputs and model weights should be placed before the forward pass. It can be one device
              for the whole module, or a dictionary mapping module name to device.
          offload (`bool`, *optional*, defaults to `False`):
              Whether or not the weights should be offloaded after the forward pass. It can be one boolean for the whole
              module, or a dictionary mapping module name to boolean.
          weights_map (`Mapping[str, torch.Tensor]`, *optional*):
              When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
          offload_buffers (`bool`, *optional*, defaults to `False`):
              Whether or not to include the associated module's buffers when offloading.
          module_name (`str`, *optional*, defaults to `""`):
              The name of the module.
          skip_keys (`str` or `List[str]`, *optional*):
              A list of keys to ignore when moving inputs or outputs between devices.
          preload_module_classes (`List[str]`, *optional*):
              A list of classes whose instances should load all their weights (even in the submodules) at the beginning
              of the forward. This should only be used for classes that have submodules which are registered but not
              called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
              `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
      """
      # If one device and one offload, we've got one hook.
      if not isinstance(execution_device, Mapping) and not isinstance(
          offload, dict
      ):
          if not offload:
              hook = AlignDevicesHook(
                  execution_device=execution_device,
                  io_same_device=True,
                  skip_keys=skip_keys,
                  place_submodules=True,
              )
              add_hook_to_module(module, hook)
          else:
              attach_align_device_hook(
                  module,
                  execution_device=execution_device,
                  offload=True,
                  weights_map=weights_map,
                  offload_buffers=offload_buffers,
                  module_name=module_name,
                  skip_keys=skip_keys,
              )
          return
  
      if not isinstance(execution_device, Mapping):
          execution_device = {
              key: execution_device for key in offload.keys()
          }
      if not isinstance(offload, Mapping):
          offload = {key: offload for key in execution_device.keys()}
  
      if (
          module_name in execution_device
          and module_name in offload
          and not offload[module_name]
      ):
          hook = AlignDevicesHook(
              execution_device=execution_device[module_name],
              offload_buffers=offload_buffers,
              io_same_device=(module_name == ""),
              place_submodules=True,
              skip_keys=skip_keys,
          )
          add_hook_to_module(module, hook)
          attach_execution_device_hook(
              module, execution_device[module_name]
          )
      elif module_name in execution_device and module_name in offload:
          attach_align_device_hook(
              module,
              execution_device=execution_device[module_name],
              offload=True,
              weights_map=weights_map,
              offload_buffers=offload_buffers,
              module_name=module_name,
              skip_keys=skip_keys,
              preload_module_classes=preload_module_classes,
          )
          if not hasattr(module, "_hf_hook"):
              hook = AlignDevicesHook(
                  execution_device=execution_device[module_name],
                  io_same_device=(module_name == ""),
                  skip_keys=skip_keys,
              )
              add_hook_to_module(module, hook)
          attach_execution_device_hook(
              module,
              execution_device[module_name],
              preload_module_classes=preload_module_classes,
              skip_keys=skip_keys,
          )
      elif module_name == "":
          hook = AlignDevicesHook(
              execution_device=execution_device.get(""),
              io_same_device=True,
              skip_keys=skip_keys,
          )
          add_hook_to_module(module, hook)
  
      for child_name, child in module.named_children():
          child_name = (
              f"{module_name}.{child_name}"
              if len(module_name) > 0
              else child_name
          )
          attach_align_device_hook_on_blocks(
              child,
              execution_device=execution_device,
              offload=offload,
              weights_map=weights_map,
              offload_buffers=offload_buffers,
              module_name=child_name,
              preload_module_classes=preload_module_classes,
              skip_keys=skip_keys,
          )
  
imports: import functools
  from typing import Dict, List, Mapping, Optional, Union
  import torch
  import torch.nn
  from .state import PartialState
  from .utils import PrefixedDataset, find_device, named_module_tensors, send_to_device, set_module_tensor_to_device
  from .utils.modeling import get_non_persistent_buffers
label_desc: Use image node_id 6 for calling a global function with example usage: attach_align_device_hook_on_blocks(module, execution_device, offload, weights_map, offload_buffers, module_name, skip_keys, preload_module_classes) without return types