code: def test_dataset_hetero_graph_split(self):
      G = generate_dense_hete_dataset()
      hete = HeteroGraph(G)
      # node
      dataset = GraphDataset([hete], task="node")
      split_res = dataset.split()
      for node_type in hete.node_label_index:
          num_nodes = int(len(hete.node_label_index[node_type]))
          node_0 = int(num_nodes * 0.8)
          node_1 = int(num_nodes * 0.1)
          node_2 = num_nodes - node_0 - node_1
  
          self.assertEqual(
              len(split_res[0][0].node_label_index[node_type]),
              node_0,
          )
  
          self.assertEqual(
              len(split_res[1][0].node_label_index[node_type]),
              node_1,
          )
  
          self.assertEqual(
              len(split_res[2][0].node_label_index[node_type]),
              node_2,
          )
  
      # node with specified split type
      dataset = GraphDataset([hete], task="node")
      node_split_types = ["n1"]
      split_res = dataset.split(split_types=node_split_types)
      for node_type in hete.node_label_index:
          if node_type in node_split_types:
              num_nodes = int(len(hete.node_label_index[node_type]))
              node_0 = int(num_nodes * 0.8)
              node_1 = int(num_nodes * 0.1)
              node_2 = num_nodes - node_0 - node_1
              self.assertEqual(
                  len(split_res[0][0].node_label_index[node_type]),
                  node_0,
              )
  
              self.assertEqual(
                  len(split_res[1][0].node_label_index[node_type]),
                  node_1,
              )
  
              self.assertEqual(
                  len(split_res[2][0].node_label_index[node_type]),
                  node_2,
              )
          else:
              num_nodes = int(len(hete.node_label_index[node_type]))
              self.assertEqual(
                  len(split_res[0][0].node_label_index[node_type]),
                  num_nodes,
              )
  
              self.assertEqual(
                  len(split_res[1][0].node_label_index[node_type]),
                  num_nodes,
              )
  
              self.assertEqual(
                  len(split_res[2][0].node_label_index[node_type]),
                  num_nodes,
              )
  
      # node with specified split type (string mode)
      dataset = GraphDataset([hete], task="node")
      node_split_types = "n1"
      split_res = dataset.split(split_types=node_split_types)
      for node_type in hete.node_label_index:
          if node_type in node_split_types:
              num_nodes = int(len(hete.node_label_index[node_type]))
              node_0 = int(num_nodes * 0.8)
              node_1 = int(num_nodes * 0.1)
              node_2 = num_nodes - node_0 - node_1
              self.assertEqual(
                  len(split_res[0][0].node_label_index[node_type]),
                  node_0,
              )
  
              self.assertEqual(
                  len(split_res[1][0].node_label_index[node_type]),
                  node_1,
              )
  
              self.assertEqual(
                  len(split_res[2][0].node_label_index[node_type]),
                  node_2,
              )
          else:
              num_nodes = int(len(hete.node_label_index[node_type]))
              self.assertEqual(
                  len(split_res[0][0].node_label_index[node_type]),
                  num_nodes,
              )
  
              self.assertEqual(
                  len(split_res[1][0].node_label_index[node_type]),
                  num_nodes,
              )
  
              self.assertEqual(
                  len(split_res[2][0].node_label_index[node_type]),
                  num_nodes,
              )
  
      # edge
      dataset = GraphDataset([hete], task="edge")
      split_res = dataset.split()
      for edge_type in hete.edge_label_index:
          num_edges = hete.edge_label_index[edge_type].shape[1]
          edge_0 = int(num_edges * 0.8)
          edge_1 = int(num_edges * 0.1)
          edge_2 = num_edges - edge_0 - edge_1
          self.assertEqual(
              split_res[0][0].edge_label_index[edge_type].shape[1],
              edge_0,
          )
  
          self.assertEqual(
              split_res[1][0].edge_label_index[edge_type].shape[1],
              edge_1,
          )
  
          self.assertEqual(
              split_res[2][0].edge_label_index[edge_type].shape[1],
              edge_2,
          )
  
      # edge with specified split type
      dataset = GraphDataset([hete], task="edge")
      edge_split_types = [("n1", "e1", "n1"), ("n1", "e2", "n2")]
      split_res = dataset.split(split_types=edge_split_types)
      for edge_type in hete.edge_label_index:
          if edge_type in edge_split_types:
              num_edges = hete.edge_label_index[edge_type].shape[1]
              edge_0 = int(num_edges * 0.8)
              edge_1 = int(num_edges * 0.1)
              edge_2 = num_edges - edge_0 - edge_1
              self.assertEqual(
                  split_res[0][0].edge_label_index[edge_type].shape[1],
                  edge_0,
              )
  
              self.assertEqual(
                  split_res[1][0].edge_label_index[edge_type].shape[1],
                  edge_1,
              )
  
              self.assertEqual(
                  split_res[2][0].edge_label_index[edge_type].shape[1],
                  edge_2,
              )
          else:
              num_edges = hete.edge_label_index[edge_type].shape[1]
              self.assertEqual(
                  split_res[0][0].edge_label_index[edge_type].shape[1],
                  num_edges,
              )
  
              self.assertEqual(
                  split_res[1][0].edge_label_index[edge_type].shape[1],
                  num_edges,
              )
  
              self.assertEqual(
                  split_res[2][0].edge_label_index[edge_type].shape[1],
                  num_edges,
              )
  
      # link_pred
      dataset = GraphDataset([hete], task="link_pred")
      split_res = dataset.split(transductive=True)
      for edge_type in hete.edge_label_index:
          num_edges = hete.edge_label_index[edge_type].shape[1]
          edge_0 = 2 * int(0.8 * num_edges)
          edge_1 = 2 * int(0.1 * num_edges)
          edge_2 = 2 * (
              num_edges - int(0.8 * num_edges) - int(0.1 * num_edges)
          )
          self.assertEqual(
              split_res[0][0].edge_label_index[edge_type].shape[1],
              edge_0,
          )
          self.assertEqual(
              split_res[1][0].edge_label_index[edge_type].shape[1],
              edge_1,
          )
          self.assertEqual(
              split_res[2][0].edge_label_index[edge_type].shape[1],
              edge_2,
          )
  
      # link_pred with specified split type
      dataset = GraphDataset([hete], task="link_pred")
      link_split_types = [("n1", "e1", "n1"), ("n1", "e2", "n2")]
      split_res = dataset.split(
          transductive=True, split_types=link_split_types
      )
  
      for edge_type in hete.edge_label_index:
          if edge_type in link_split_types:
              num_edges = hete.edge_label_index[edge_type].shape[1]
              edge_0 = 2 * int(0.8 * num_edges)
              edge_1 = 2 * int(0.1 * num_edges)
              edge_2 = 2 * (
                  num_edges
                  - int(0.8 * num_edges)
                  - int(0.1 * num_edges)
              )
              self.assertEqual(
                  split_res[0][0].edge_label_index[edge_type].shape[1],
                  edge_0,
              )
              self.assertEqual(
                  split_res[1][0].edge_label_index[edge_type].shape[1],
                  edge_1,
              )
              self.assertEqual(
                  split_res[2][0].edge_label_index[edge_type].shape[1],
                  edge_2,
              )
          else:
              num_edges = hete.edge_label_index[edge_type].shape[1]
              self.assertEqual(
                  split_res[0][0].edge_label_index[edge_type].shape[1],
                  num_edges,
              )
              self.assertEqual(
                  split_res[1][0].edge_label_index[edge_type].shape[1],
                  num_edges,
              )
              self.assertEqual(
                  split_res[2][0].edge_label_index[edge_type].shape[1],
                  num_edges,
              )
  
      # link_pred + disjoint
      dataset = GraphDataset(
          [hete],
          task="link_pred",
          edge_train_mode="disjoint",
          edge_message_ratio=0.5,
      )
      split_res = dataset.split(
          transductive=True,
          split_ratio=[0.6, 0.2, 0.2],
      )
      for edge_type in hete.edge_label_index:
          num_edges = hete.edge_label_index[edge_type].shape[1]
          edge_0 = int(0.6 * num_edges)
          edge_0 = 2 * (edge_0 - int(0.5 * edge_0))
          edge_1 = 2 * int(0.2 * num_edges)
          edge_2 = 2 * (
              num_edges - int(0.6 * num_edges) - int(0.2 * num_edges)
          )
  
          self.assertEqual(
              split_res[0][0].edge_label_index[edge_type].shape[1],
              edge_0,
          )
          self.assertEqual(
              split_res[1][0].edge_label_index[edge_type].shape[1],
              edge_1,
          )
          self.assertEqual(
              split_res[2][0].edge_label_index[edge_type].shape[1],
              edge_2,
          )
  
      # link pred with edge_split_mode set to "exact"
      dataset = GraphDataset(
          [hete], task="link_pred", edge_split_mode="approximate"
      )
      split_res = dataset.split(transductive=True)
      hete_link_train_edge_num = 0
      hete_link_test_edge_num = 0
      hete_link_val_edge_num = 0
      num_edges = 0
      for edge_type in hete.edge_label_index:
          num_edges += hete.edge_label_index[edge_type].shape[1]
          if edge_type in split_res[0][0].edge_label_index:
              hete_link_train_edge_num += (
                  split_res[0][0].edge_label_index[edge_type].shape[1]
              )
          if edge_type in split_res[1][0].edge_label_index:
              hete_link_test_edge_num += (
                  split_res[1][0].edge_label_index[edge_type].shape[1]
              )
          if edge_type in split_res[2][0].edge_label_index:
              hete_link_val_edge_num += (
                  split_res[2][0].edge_label_index[edge_type].shape[1]
              )
  
      edge_0 = 2 * int(0.8 * num_edges)
      edge_1 = 2 * int(0.1 * num_edges)
      edge_2 = 2 * (
          num_edges - int(0.8 * num_edges) - int(0.1 * num_edges)
      )
      self.assertEqual(hete_link_train_edge_num, edge_0)
      self.assertEqual(hete_link_test_edge_num, edge_1)
      self.assertEqual(hete_link_val_edge_num, edge_2)
  
      # link pred with specified types and edge_split_mode set to "exact"
      dataset = GraphDataset(
          [hete],
          task="link_pred",
          edge_split_mode="approximate",
      )
      link_split_types = [("n1", "e1", "n1"), ("n1", "e2", "n2")]
      split_res = dataset.split(
          transductive=True,
          split_types=link_split_types,
      )
      hete_link_train_edge_num = 0
      hete_link_test_edge_num = 0
      hete_link_val_edge_num = 0
  
      num_split_type_edges = 0
      num_non_split_type_edges = 0
      for edge_type in hete.edge_label_index:
          if edge_type in link_split_types:
              num_split_type_edges += hete.edge_label_index[
                  edge_type
              ].shape[1]
          else:
              num_non_split_type_edges += hete.edge_label_index[
                  edge_type
              ].shape[1]
          if edge_type in split_res[0][0].edge_label_index:
              hete_link_train_edge_num += (
                  split_res[0][0].edge_label_index[edge_type].shape[1]
              )
          if edge_type in split_res[1][0].edge_label_index:
              hete_link_test_edge_num += (
                  split_res[1][0].edge_label_index[edge_type].shape[1]
              )
          if edge_type in split_res[2][0].edge_label_index:
              hete_link_val_edge_num += (
                  split_res[2][0].edge_label_index[edge_type].shape[1]
              )
  
      num_edges = num_split_type_edges
      edge_0 = 2 * int(0.8 * num_edges) + num_non_split_type_edges
      edge_1 = 2 * int(0.1 * num_edges) + num_non_split_type_edges
      edge_2 = (
          2 * (num_edges - int(0.8 * num_edges) - int(0.1 * num_edges))
          + num_non_split_type_edges
      )
  
      self.assertEqual(hete_link_train_edge_num, edge_0)
      self.assertEqual(hete_link_test_edge_num, edge_1)
      self.assertEqual(hete_link_val_edge_num, edge_2)
  
imports: import random
  import torch
  import unittest
  from torch_geometric.datasets import TUDataset, Planetoid
  from copy import deepcopy
  from deepsnap.graph import Graph
  from deepsnap.hetero_graph import HeteroGraph
  from deepsnap.dataset import GraphDataset, Generator, EnsembleGenerator
  from tests.utils import simple_networkx_graph, simple_networkx_small_graph, simple_networkx_graph_alphabet, simple_networkx_dense_graph, simple_networkx_dense_multigraph, simple_networkx_multigraph, generate_dense_hete_dataset, generate_simple_small_hete_graph, generate_simple_dense_hete_graph, generate_simple_dense_hete_multigraph, generate_dense_hete_multigraph, gen_graph
label_desc: Use image node_id 3 for calling the TestDataset obj's underlying member method code with example usage: obj.test_dataset_hetero_graph_split() without return types