code: def test_to_captum_input(mask_type):
      num_nodes = x.shape[0]
      num_node_feats = x.shape[1]
      num_edges = edge_index.shape[1]
  
      # Check for Data:
      data = Data(x, edge_index)
      args = "test_args"
      inputs, additional_forward_args = to_captum_input(
          data.x, data.edge_index, mask_type, args
      )
      if mask_type == "node":
          assert len(inputs) == 1
          assert inputs[0].shape == (1, num_nodes, num_node_feats)
          assert len(additional_forward_args) == 2
          assert torch.allclose(additional_forward_args[0], edge_index)
      elif mask_type == "edge":
          assert len(inputs) == 1
          assert inputs[0].shape == (1, num_edges)
          assert inputs[0].sum() == num_edges
          assert len(additional_forward_args) == 3
          assert torch.allclose(additional_forward_args[0], x)
          assert torch.allclose(additional_forward_args[1], edge_index)
      else:
          assert len(inputs) == 2
          assert inputs[0].shape == (1, num_nodes, num_node_feats)
          assert inputs[1].shape == (1, num_edges)
          assert inputs[1].sum() == num_edges
          assert len(additional_forward_args) == 2
          assert torch.allclose(additional_forward_args[0], edge_index)
  
      # Check for HeteroData:
      data = HeteroData()
      x2 = torch.rand(8, 3)
      data["paper"].x = x
      data["author"].x = x2
      data["paper", "to", "author"].edge_index = edge_index
      data["author", "to", "paper"].edge_index = edge_index.flip([0])
      inputs, additional_forward_args = to_captum_input(
          data.x_dict, data.edge_index_dict, mask_type, args
      )
      if mask_type == "node":
          assert len(inputs) == 2
          assert inputs[0].shape == (1, num_nodes, num_node_feats)
          assert inputs[1].shape == (1, num_nodes, num_node_feats)
          assert len(additional_forward_args) == 2
          for key in data.edge_types:
              torch.allclose(
                  additional_forward_args[0][key], data[key].edge_index
              )
      elif mask_type == "edge":
          assert len(inputs) == 2
          assert inputs[0].shape == (1, num_edges)
          assert inputs[1].shape == (1, num_edges)
          assert inputs[1].sum() == inputs[0].sum() == num_edges
          assert len(additional_forward_args) == 3
          for key in data.node_types:
              torch.allclose(
                  additional_forward_args[0][key], data[key].x
              )
          for key in data.edge_types:
              torch.allclose(
                  additional_forward_args[1][key], data[key].edge_index
              )
      else:
          assert len(inputs) == 4
          assert inputs[0].shape == (1, num_nodes, num_node_feats)
          assert inputs[1].shape == (1, num_nodes, num_node_feats)
          assert inputs[2].shape == (1, num_edges)
          assert inputs[3].shape == (1, num_edges)
          assert inputs[3].sum() == inputs[2].sum() == num_edges
          assert len(additional_forward_args) == 2
          for key in data.edge_types:
              torch.allclose(
                  additional_forward_args[0][key], data[key].edge_index
              )
  
file_path: pytorch_geometric/test/explain/algorithm/test_captum.py
imports: import pytest
  import torch
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.explain.algorithm.captum import to_captum_input
  from torch_geometric.nn import GAT, GCN, SAGEConv
  from torch_geometric.nn.conv import MessagePassing
  from torch_geometric.nn.models import to_captum_model
  from torch_geometric.testing import withPackage
label_desc: Use image node_id 4 for calling a global function with example usage: test_to_captum_input(mask_type) without return types