code: def test_hetero_data_subgraph():
      data = HeteroData()
      data.num_node_types = 3
      data["paper"].x = x_paper
      data["paper"].name = "paper"
      data["paper"].num_nodes = x_paper.size(0)
      data["author"].x = x_author
      data["author"].num_nodes = x_author.size(0)
      data["conf"].x = x_conference
      data["conf"].num_nodes = x_conference.size(0)
      data["paper", "paper"].edge_index = edge_index_paper_paper
      data["paper", "paper"].edge_attr = edge_attr_paper_paper
      data["paper", "paper"].name = "cites"
      data["author", "paper"].edge_index = edge_index_author_paper
      data["paper", "author"].edge_index = edge_index_paper_author
      data["paper", "conf"].edge_index = edge_index_paper_conference
  
      subset = {
          "paper": torch.randperm(x_paper.size(0))[:4],
          "author": torch.randperm(x_author.size(0))[:2],
          "conf": torch.randperm(x_conference.size(0))[:2],
      }
  
      out = data.subgraph(subset)
      out.validate(raise_on_error=True)
  
      assert out.num_node_types == data.num_node_types
      assert out.node_types == ["paper", "author", "conf"]
  
      for key in out.node_types:
          assert len(out[key]) == len(data[key])
          assert torch.allclose(out[key].x, data[key].x[subset[key]])
          assert out[key].num_nodes == subset[key].size(0)
          if key == "paper":
              assert out["paper"].name == "paper"
  
      # Construct correct edge index manually:
      node_mask = (
          {}
      )  # for each node type a mask of nodes in the subgraph
      node_map = (
          {}
      )  # for each node type a map from old node id to new node id
      for key in out.node_types:
          node_mask[key] = torch.zeros(
              (data[key].num_nodes,), dtype=torch.bool
          )
          node_map[key] = torch.zeros(
              (data[key].num_nodes,), dtype=torch.long
          )
          node_mask[key][subset[key]] = True
          node_map[key][subset[key]] = torch.arange(subset[key].size(0))
  
      edge_mask = (
          {}
      )  # for each edge type a mask of edges in the subgraph
      subgraph_edge_index = (
          {}
      )  # for each edge type the edge index of the subgraph
      for key in out.edge_types:
          edge_mask[key] = (
              node_mask[key[0]][data[key].edge_index[0]]
              & node_mask[key[-1]][data[key].edge_index[1]]
          )
          subgraph_edge_index[key] = data[key].edge_index[
              :, edge_mask[key]
          ]
          subgraph_edge_index[key][0] = node_map[key[0]][
              subgraph_edge_index[key][0]
          ]
          subgraph_edge_index[key][1] = node_map[key[-1]][
              subgraph_edge_index[key][1]
          ]
  
      assert out.edge_types == [
          ("paper", "to", "paper"),
          ("author", "to", "paper"),
          ("paper", "to", "author"),
          ("paper", "to", "conf"),
      ]
  
      for key in out.edge_types:
          assert len(out[key]) == len(data[key])
          assert torch.equal(
              out[key].edge_index, subgraph_edge_index[key]
          )
          if key == ("paper", "to", "paper"):
              assert torch.allclose(
                  out[key].edge_attr,
                  data[key].edge_attr[edge_mask[key]],
              )
              assert out[key].name == "cites"
  
      # Test for bool and long in `subset_dict`.
      author_mask = torch.zeros((x_author.size(0),), dtype=torch.bool)
      author_mask[subset["author"]] = True
      subset_mixed = {
          "paper": subset["paper"],
          "author": author_mask,
      }
      out = data.subgraph(subset_mixed)
      out.validate(raise_on_error=True)
  
      assert out.num_node_types == data.num_node_types
      assert out.node_types == ["paper", "author", "conf"]
      assert out["paper"].num_nodes == subset["paper"].size(0)
      assert out["author"].num_nodes == subset["author"].size(0)
      assert out["conf"].num_nodes == data["conf"].num_nodes
      assert out.edge_types == [
          ("paper", "to", "paper"),
          ("author", "to", "paper"),
          ("paper", "to", "author"),
          ("paper", "to", "conf"),
      ]
  
      out = data.node_type_subgraph(["paper", "author"])
      assert out.node_types == ["paper", "author"]
      assert out.edge_types == [
          ("paper", "to", "paper"),
          ("author", "to", "paper"),
          ("paper", "to", "author"),
      ]
  
      out = data.edge_type_subgraph([("paper", "author")])
      assert out.node_types == ["paper", "author"]
      assert out.edge_types == [("paper", "to", "author")]
  
      subset = {
          ("paper", "to", "paper"): torch.arange(4),
      }
  
      out = data.edge_subgraph(subset)
      assert out.node_types == data.node_types
      assert out.edge_types == data.edge_types
      assert data["paper"] == out["paper"]
      assert data["author"] == out["author"]
      assert data["paper", "author"] == out["paper", "author"]
      assert data["author", "paper"] == out["author", "paper"]
  
      assert out["paper", "paper"].num_edges == 4
      assert out["paper", "paper"].edge_index.size() == (2, 4)
      assert out["paper", "paper"].edge_attr.size() == (4, 8)
  
imports: import copy
  import pytest
  import torch
  from torch_geometric.data import HeteroData
  from torch_geometric.data.storage import EdgeStorage
  from torch_geometric.testing import get_random_edge_index, get_random_tensor_frame, withPackage
  from torch_geometric.typing import TensorFrame
label_desc: Use image node_id 7 for calling a global function with example usage: test_hetero_data_subgraph() without return types