code: def test_dataset_split(self):
      # inductively split with graph task
      pyg_dataset = TUDataset("./enzymes", "ENZYMES")
      graphs = GraphDataset.pyg_to_graphs(pyg_dataset)
      dataset = GraphDataset(graphs, task="graph")
      split_res = dataset.split(transductive=False)
      num_graphs = len(dataset)
      num_train = int(0.8 * num_graphs)
      num_val = int(0.1 * num_graphs)
      num_test = num_graphs - num_train - num_val
      self.assertEqual(num_train, len(split_res[0]))
      self.assertEqual(num_val, len(split_res[1]))
      self.assertEqual(num_test, len(split_res[2]))
  
      # inductively split with link_pred task
      # and default (`all`) edge_train_mode
      pyg_dataset = TUDataset("./enzymes", "ENZYMES")
      graphs = GraphDataset.pyg_to_graphs(pyg_dataset)
      dataset = GraphDataset(graphs, task="link_pred")
      split_res = dataset.split(transductive=False)
      num_graphs = len(dataset)
      num_train = int(num_graphs * 0.8)
      num_val = int(num_graphs * 0.1)
      num_test = num_graphs - num_train - num_val
      self.assertEqual(num_train, len(split_res[0]))
      self.assertEqual(num_val, len(split_res[1]))
      self.assertEqual(num_test, len(split_res[2]))
  
      # inductively split with link_pred task and `disjoint` edge_train_mode
      pyg_dataset = TUDataset("./enzymes", "ENZYMES")
      graphs = GraphDataset.pyg_to_graphs(pyg_dataset)
      dataset = GraphDataset(
          graphs,
          task="link_pred",
          edge_train_mode="disjoint",
      )
      split_res = dataset.split(transductive=False)
      num_graphs = len(dataset)
      num_train = int(num_graphs * 0.8)
      num_val = int(num_graphs * 0.1)
      num_test = num_graphs - num_train - num_val
      self.assertEqual(num_train, len(split_res[0]))
      self.assertEqual(num_val, len(split_res[1]))
      self.assertEqual(num_test, len(split_res[2]))
  
      # transductively split with node task
      pyg_dataset = Planetoid("./cora", "Cora")
      graphs = GraphDataset.pyg_to_graphs(pyg_dataset)
      dataset = GraphDataset(graphs, task="node")
      num_nodes = dataset.num_nodes[0]
      node_0 = int(0.8 * num_nodes)
      node_1 = int(0.1 * num_nodes)
      node_2 = num_nodes - node_0 - node_1
      split_res = dataset.split()
      self.assertEqual(len(split_res[0][0].node_label_index), node_0)
      self.assertEqual(len(split_res[1][0].node_label_index), node_1)
      self.assertEqual(len(split_res[2][0].node_label_index), node_2)
  
      for j in range(3):
          for i in range(split_res[j][0].node_label_index.shape[0]):
              node = split_res[j][0].node_label_index[i].item()
              node_label = split_res[j][0].node_label[i].item()
              self.assertEqual(
                  dataset[0].G.nodes[node]["node_label"], node_label
              )
  
      # transductively split with edge task
      (
          G,
          x,
          y,
          edge_x,
          edge_y,
          edge_index,
          graph_x,
          graph_y,
      ) = simple_networkx_graph()
      Graph.add_edge_attr(G, "edge_feature", edge_x)
      Graph.add_edge_attr(G, "edge_label", edge_y)
      Graph.add_node_attr(G, "node_feature", x)
      Graph.add_node_attr(G, "node_label", y)
      Graph.add_graph_attr(G, "graph_feature", graph_x)
      Graph.add_graph_attr(G, "graph_label", graph_y)
      graph = Graph(G)
      num_edges = graph.num_edges
      graphs = [graph]
      dataset = GraphDataset(graphs, task="edge")
      split_res = dataset.split()
      edge_0 = int(0.8 * num_edges)
      self.assertEqual(
          split_res[0][0].edge_label_index.shape[1],
          edge_0,
      )
      edge_1 = int(0.1 * num_edges)
      self.assertEqual(
          split_res[1][0].edge_label_index.shape[1],
          edge_1,
      )
      self.assertEqual(
          split_res[2][0].edge_label_index.shape[1],
          num_edges - edge_0 - edge_1,
      )
  
      for j in range(3):
          for i in range(split_res[j][0].edge_label_index.shape[1]):
              node_0 = split_res[j][0].edge_label_index[0][i].item()
              node_1 = split_res[j][0].edge_label_index[1][i].item()
              edge_label = split_res[j][0].edge_label[i].item()
              self.assertEqual(
                  G.edges[node_0, node_1]["edge_label"], edge_label
              )
  
      # transductively split with link_pred task
      # and default (`all`) edge_train_mode
      pyg_dataset = Planetoid("./cora", "Cora")
      # dataset is undirected
      graphs = GraphDataset.pyg_to_graphs(pyg_dataset)
      dataset = GraphDataset(graphs, task="link_pred")
      num_edges = dataset.num_edges[0]
      edge_0 = 2 * 2 * int(0.8 * num_edges)
      edge_1 = 2 * 2 * int(0.1 * num_edges)
      edge_2 = (
          2
          * 2
          * (num_edges - int(0.8 * num_edges) - int(0.1 * num_edges))
      )
      split_res = dataset.split()
      self.assertEqual(
          split_res[0][0].edge_label_index.shape[1], edge_0
      )
      self.assertEqual(
          split_res[1][0].edge_label_index.shape[1], edge_1
      )
      self.assertEqual(
          split_res[2][0].edge_label_index.shape[1], edge_2
      )
  
      # transductively split with link_pred task, `split` edge_train_mode
      # and 0.5 edge_message_ratio
      pyg_dataset = Planetoid("./cora", "Cora")
      graphs = GraphDataset.pyg_to_graphs(pyg_dataset)
      dataset = GraphDataset(
          graphs,
          task="link_pred",
          edge_train_mode="disjoint",
          edge_message_ratio=0.5,
      )
      num_edges = dataset.num_edges[0]
      split_res = dataset.split()
      edge_0 = 2 * int(0.8 * num_edges)
      edge_0 = 2 * (edge_0 - int(0.5 * edge_0))
      edge_1 = 2 * 2 * int(0.1 * num_edges)
      edge_2 = (
          2
          * 2
          * (num_edges - int(0.8 * num_edges) - int(0.1 * num_edges))
      )
      self.assertEqual(
          split_res[0][0].edge_label_index.shape[1], edge_0
      )
      self.assertEqual(
          split_res[1][0].edge_label_index.shape[1], edge_1
      )
      self.assertEqual(
          split_res[2][0].edge_label_index.shape[1], edge_2
      )
  
      # resample disjoint
      self.assertEqual(
          split_res[0][0].edge_label_index.shape[1], edge_0
      )
      self.assertEqual(
          split_res[1][0].edge_label_index.shape[1], edge_1
      )
      self.assertEqual(
          split_res[2][0].edge_label_index.shape[1], edge_2
      )
  
      # transductively split with link_pred task
      # and specified edge_negative_sampling_ratio
      pyg_dataset = Planetoid("./cora", "Cora")
      graphs = GraphDataset.pyg_to_graphs(pyg_dataset)
      dataset = GraphDataset(
          graphs, task="link_pred", edge_negative_sampling_ratio=2
      )
      num_edges = dataset.num_edges[0]
      edge_0 = (2 + 1) * 2 * int(0.8 * num_edges)
      edge_1 = (2 + 1) * 2 * int(0.1 * num_edges)
      edge_2 = (
          (2 + 1)
          * 2
          * (num_edges - int(0.8 * num_edges) - int(0.1 * num_edges))
      )
      split_res = dataset.split()
      self.assertEqual(
          split_res[0][0].edge_label_index.shape[1], edge_0
      )
      self.assertEqual(
          split_res[1][0].edge_label_index.shape[1], edge_1
      )
      self.assertEqual(
          split_res[2][0].edge_label_index.shape[1], 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 4 for calling the TestDataset obj's underlying member method code with example usage: obj.test_dataset_split() without return types