code: def _generate_align_sparsity(
      masks: _MASKS, target_spaces: _PRUNING_TARGET_SPACES
  ) -> _MASKS:
      align_masks = defaultdict(dict)
      for module_name, ts in target_spaces.items():
          for target_name, target_space in ts.items():
              align_module_name = target_space.align["module_name"]  # type: ignore
              align_module_name = (
                  module_name
                  if align_module_name is None
                  else align_module_name
              )
              assert align_module_name in masks
              src_mask = masks[align_module_name][target_space.align["target_name"]]  # type: ignore
              align_dims: List[int] = target_space.align["dims"]  # type: ignore
              reduce_dims = [
                  d
                  for d in range(len(src_mask.shape))
                  if d not in align_dims
                  and d - len(src_mask.shape) not in align_dims
              ]
              align_mask = src_mask.sum(reduce_dims).bool().float()
              if target_space._scaler is not None:
                  assert (
                      target_space.shape is not None
                  ), f"The shape of {module_name}.{target_name} is not tracked"
                  align_mask = (
                      target_space._scaler.shrink(
                          target_space._scaler.expand(
                              align_mask, target_space.shape
                          ),
                          keepdim=True,
                      )
                      .bool()
                      .float()
                  )
              align_masks[module_name][target_name] = align_mask
      return align_masks
  
imports: from __future__ import annotations
  from collections import defaultdict
  from functools import reduce
  import heapq
  from typing import Callable, Dict, List, Tuple
  import numpy
  import torch
  from .calculate_metrics import _METRICS
  from ...base.compressor import _PRUNING_TARGET_SPACES
  from ...base.target_space import PruningTargetSpace, TargetType
label_desc: Use image node_id 5 for calling a global function with example usage: _generate_align_sparsity(masks, target_spaces) and returns: align_masks