code: def to_captum_input(
      x: Union[Tensor, Dict[NodeType, Tensor]],
      edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
      mask_type: Union[str, MaskLevelType],
      *args,
  ) -> Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...]]:
      r"""Given :obj:`x`, :obj:`edge_index` and :obj:`mask_type`, converts it
      to a format to use in `Captum <https://captum.ai/>`_ attribution
      methods. Returns :obj:`inputs` and :obj:`additional_forward_args`
      required for :captum:`Captum's` :obj:`attribute` functions.
      See :meth:`~torch_geometric.nn.models.to_captum_model` for example usage.
  
      Args:
          x (torch.Tensor or Dict[NodeType, torch.Tensor]): The node features.
              For heterogeneous graphs this is a dictionary holding node featues
              for each node type.
          edge_index(torch.Tensor or Dict[EdgeType, torch.Tensor]): The edge
              indices. For heterogeneous graphs this is a dictionary holding the
              :obj:`edge index` for each edge type.
          mask_type (str): Denotes the type of mask to be created with
              a Captum explainer. Valid inputs are :obj:`"edge"`, :obj:`"node"`,
              and :obj:`"node_and_edge"`.
          *args: Additional forward arguments of the model being explained
              which will be added to :obj:`additional_forward_args`.
      """
      mask_type = MaskLevelType(mask_type)
  
      additional_forward_args = []
      if isinstance(x, Tensor) and isinstance(edge_index, Tensor):
          if mask_type == MaskLevelType.node:
              inputs = [x.unsqueeze(0)]
          elif mask_type == MaskLevelType.edge:
              inputs = [_to_edge_mask(edge_index).unsqueeze(0)]
              additional_forward_args.append(x)
          else:
              inputs = [
                  x.unsqueeze(0),
                  _to_edge_mask(edge_index).unsqueeze(0),
              ]
          additional_forward_args.append(edge_index)
  
      elif isinstance(x, Dict) and isinstance(edge_index, Dict):
          node_types = x.keys()
          edge_types = edge_index.keys()
          inputs = []
          if mask_type == MaskLevelType.node:
              for key in node_types:
                  inputs.append(x[key].unsqueeze(0))
          elif mask_type == MaskLevelType.edge:
              for key in edge_types:
                  inputs.append(
                      _to_edge_mask(edge_index[key]).unsqueeze(0)
                  )
              additional_forward_args.append(x)
          else:
              for key in node_types:
                  inputs.append(x[key].unsqueeze(0))
              for key in edge_types:
                  inputs.append(
                      _to_edge_mask(edge_index[key]).unsqueeze(0)
                  )
          additional_forward_args.append(edge_index)
  
      else:
          raise ValueError(
              "'x' and 'edge_index' need to be either"
              f"'Dict' or 'Tensor' got({type(x)}, {type(edge_index)})"
          )
  
      additional_forward_args.extend(args)
  
      return tuple(inputs), tuple(additional_forward_args)
  
imports: from enum import Enum
  from typing import Dict, Optional, Tuple, Union
  import torch
  from torch import Tensor
  from torch_geometric.explain.algorithm.utils import clear_masks, set_hetero_masks, set_masks
  from torch_geometric.explain.config import ModelConfig, ModelMode, ModelReturnType
  from torch_geometric.typing import EdgeType, Metadata, NodeType
label_desc: Use image node_id 2 for calling a global function with example usage: to_captum_input(x, edge_index, mask_type) and returns: tuple,  tuple