code: def _sort_key_list(
      keys: List[device], device_ids: Union[None, List[int]] = None
  ) -> List[device]:
      """
      Sorts list of torch devices (keys) by given index list, device_ids. If keys
      contains only one device, then the list is returned unchanged. If keys
      contains a device for which the id is not contained in device_ids, then
      an error is returned. This method is used to identify the order of DataParallel
      batched devices, given the device ID ordering.
      """
      if len(keys) == 1:
          return keys
      id_dict: Dict[int, device] = {}
      assert (
          device_ids is not None
      ), "Device IDs must be provided with multiple devices."
      for key in keys:
          if key.index in id_dict:
              raise AssertionError(
                  "Duplicate CUDA Device ID identified in device list."
              )
          id_dict[key.index] = key
  
      out_list = [
          id_dict[device_id]
          for device_id in filter(
              lambda device_id: device_id in id_dict, device_ids
          )
      ]
  
      assert len(out_list) == len(
          keys
      ), "Given Device ID List does not match"
      "devices with computed tensors."
  
      return out_list
  
imports: import typing
  from enum import Enum
  from functools import reduce
  from inspect import signature
  from typing import Any, Callable, cast, Dict, List, overload, Tuple, Union
  import numpy
  import torch
  from captum._utils.typing import BaselineType, Literal, TargetType, TensorOrTupleOfTensorsGeneric, TupleOrTensorOrBoolGeneric
  from torch import device, Tensor
  from torch.nn import Module
label_desc: Use image node_id 43 for calling a global function with example usage: _sort_key_list(keys, device_ids) and returns: out_list, keys