code: def _expand_masks(
      masks: _MASKS, target_spaces: _PRUNING_TARGET_SPACES
  ) -> _MASKS:
      # expand the mask shape from metric shape to target shape
      new_masks = defaultdict(dict)
      for module_name, module_masks in masks.items():
          for target_name, target_mask in module_masks.items():
              target_space = target_spaces[module_name][target_name]
              if target_space._scaler:
                  new_masks[module_name][
                      target_name
                  ] = target_space._scaler.expand(
                      target_mask,
                      target_space.shape,
                      keepdim=True,
                      full_expand=False,
                  )  # type: ignore
              else:
                  new_masks[module_name][target_name] = target_mask
      return new_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 13 for calling a global function with example usage: _expand_masks(masks, target_spaces) and returns: new_masks