code: def _extract_device_ids(
      forward_fn: Callable,
      saved_layer: Dict[Module, Dict[device, Tuple[Tensor, ...]]],
      device_ids: Union[None, List[int]],
  ) -> Union[None, List[int]]:
      r"""
      A helper function to extract device_ids from `forward_function` in case it is
      provided as part of a `DataParallel` model or if is accessible from
      `forward_fn`.
      In case input device_ids is not None, this function returns that value.
      """
      # Multiple devices / keys implies a DataParallel model, so we look for
      # device IDs if given or available from forward function
      # (DataParallel model object).
      if (
          max(
              len(saved_layer[single_layer])
              for single_layer in saved_layer
          )
          > 1
          and device_ids is None
      ):
          if (
              hasattr(forward_fn, "device_ids")
              and cast(Any, forward_fn).device_ids is not None
          ):
              device_ids = cast(Any, forward_fn).device_ids
          else:
              raise AssertionError(
                  "Layer tensors are saved on multiple devices, however unable to access"
                  " device ID list from the `forward_fn`. Device ID list must be"
                  " accessible from `forward_fn`. For example, they can be retrieved"
                  " if `forward_fn` is a model of type `DataParallel`. It is used"
                  " for identifying device batch ordering."
              )
      return device_ids
  
imports: import threading
  import typing
  import warnings
  from collections import defaultdict
  from typing import Any, Callable, cast, Dict, List, Optional, Sequence, Tuple, Union
  import torch
  from captum._utils.common import _reduce_list, _run_forward, _sort_key_list, _verify_select_neuron
  from captum._utils.sample_gradient import SampleGradientWrapper
  from captum._utils.typing import Literal, ModuleOrModuleList, TargetType, TensorOrTupleOfTensorsGeneric
  from torch import device, Tensor
  from torch.nn import Module
label_desc: Use image node_id 12 for calling a global function with example usage: _extract_device_ids(forward_fn, saved_layer, device_ids) and returns: device_ids