code: def test_dataset_split_custom(self):
      # transductive split with node task (self defined dataset)
      (
          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)
  
      num_nodes = len(list(G.nodes))
      nodes_train = torch.tensor(list(G.nodes)[: int(0.3 * num_nodes)])
      nodes_val = torch.tensor(
          list(G.nodes)[int(0.3 * num_nodes) : int(0.6 * num_nodes)]
      )
      nodes_test = torch.tensor(list(G.nodes)[int(0.6 * num_nodes) :])
  
      graph_train = Graph(
          node_feature=x,
          node_label=y,
          edge_index=edge_index,
          node_label_index=nodes_train,
          directed=True,
      )
      graph_val = Graph(
          node_feature=x,
          node_label=y,
          edge_index=edge_index,
          node_label_index=nodes_val,
          directed=True,
      )
      graph_test = Graph(
          node_feature=x,
          node_label=y,
          edge_index=edge_index,
          node_label_index=nodes_test,
          directed=True,
      )
  
      graphs_train = [graph_train]
      graphs_val = [graph_val]
      graphs_test = [graph_test]
  
      dataset_train, dataset_val, dataset_test = (
          GraphDataset(graphs_train, task="node"),
          GraphDataset(graphs_val, task="node"),
          GraphDataset(graphs_test, task="node"),
      )
  
      self.assertEqual(
          dataset_train[0].node_label_index.tolist(),
          list(range(int(0.3 * num_nodes))),
      )
      self.assertEqual(
          dataset_val[0].node_label_index.tolist(),
          list(range(int(0.3 * num_nodes), int(0.6 * num_nodes))),
      )
      self.assertEqual(
          dataset_test[0].node_label_index.tolist(),
          list(range(int(0.6 * num_nodes), num_nodes)),
      )
  
      # transductive split with link_pred task (train/val split)
      edges = list(G.edges)
      num_edges = len(edges)
      edges_train = edges[: int(0.7 * num_edges)]
      edges_val = edges[int(0.7 * num_edges) :]
      link_size_list = [len(edges_train), len(edges_val)]
  
      # generate pseudo pos and neg edges, they may overlap here
      train_pos = torch.LongTensor(edges_train).permute(1, 0)
      val_pos = torch.LongTensor(edges_val).permute(1, 0)
      val_neg = torch.randint(
          high=10, size=val_pos.shape, dtype=torch.int64
      )
      val_neg_double = torch.cat((val_neg, val_neg), dim=1)
  
      num_train = len(edges_train)
      num_val = len(edges_val)
  
      graph_train = Graph(
          node_feature=x,
          edge_index=edge_index,
          edge_feature=edge_x,
          directed=True,
          edge_label_index=train_pos,
      )
  
      graph_val = Graph(
          node_feature=x,
          edge_index=edge_index,
          edge_feature=edge_x,
          directed=True,
          edge_label_index=val_pos,
          negative_edge=val_neg_double,
      )
  
      graphs_train = [graph_train]
      graphs_val = [graph_val]
  
      dataset_train, dataset_val = (
          GraphDataset(
              graphs_train, task="link_pred", resample_negatives=True
          ),
          GraphDataset(
              graphs_val,
              task="link_pred",
              edge_negative_sampling_ratio=2,
          ),
      )
  
      self.assertEqual(
          dataset_train[0].edge_label_index.shape[1],
          2 * link_size_list[0],
      )
      self.assertEqual(
          dataset_train[0].edge_label.shape[0], 2 * link_size_list[0]
      )
      self.assertEqual(
          dataset_val[0].edge_label_index.shape[1],
          val_pos.shape[1] + val_neg_double.shape[1],
      )
      self.assertEqual(
          dataset_val[0].edge_label.shape[0],
          val_pos.shape[1] + val_neg_double.shape[1],
      )
      self.assertTrue(
          torch.equal(
              dataset_train[0].edge_label_index[:, :num_train],
              train_pos,
          )
      )
      self.assertTrue(
          torch.equal(
              dataset_val[0].edge_label_index[:, :num_val], val_pos
          )
      )
      self.assertTrue(
          torch.equal(
              dataset_val[0].edge_label_index[:, num_val:],
              val_neg_double,
          )
      )
  
      dataset_train.resample_negatives = False
      self.assertTrue(
          torch.equal(
              dataset_train[0].edge_label_index,
              dataset_train[0].edge_label_index,
          )
      )
  
      # transductive split with link_pred task with edge label
      edge_label_train = torch.LongTensor(
          [1, 2, 3, 2, 1, 1, 2, 3, 2, 0, 0]
      )
      edge_label_val = torch.LongTensor([1, 2, 3, 2, 1, 0])
  
      graph_train = Graph(
          node_feature=x,
          edge_index=edge_index,
          directed=True,
          edge_label_index=train_pos,
          edge_label=edge_label_train,
      )
  
      graph_val = Graph(
          node_feature=x,
          edge_index=edge_index,
          directed=True,
          edge_label_index=val_pos,
          negative_edge=val_neg,
          edge_label=edge_label_val,
      )
  
      graphs_train = [graph_train]
      graphs_val = [graph_val]
  
      dataset_train, dataset_val = (
          GraphDataset(graphs_train, task="link_pred"),
          GraphDataset(graphs_val, task="link_pred"),
      )
  
      self.assertTrue(
          torch.equal(
              dataset_train[0].edge_label_index,
              dataset_train[0].edge_label_index,
          )
      )
  
      self.assertTrue(
          torch.equal(
              dataset_train[0].edge_label[:num_train], edge_label_train
          )
      )
  
      self.assertTrue(
          torch.equal(
              dataset_val[0].edge_label[:num_val], edge_label_val
          )
      )
  
      # Multiple graph tensor backend link prediction (inductive)
      pyg_dataset = Planetoid("./cora", "Cora")
      x = pyg_dataset[0].x
      y = pyg_dataset[0].y
      edge_index = pyg_dataset[0].edge_index
      row, col = edge_index
      mask = row < col
      row, col = row[mask], col[mask]
      edge_index = torch.stack([row, col], dim=0)
      edge_index = torch.cat(
          [edge_index, torch.flip(edge_index, [0])], dim=1
      )
  
      graphs = [
          Graph(
              node_feature=x,
              node_label=y,
              edge_index=edge_index,
              directed=False,
          )
      ]
      graphs = [copy.deepcopy(graphs[0]) for _ in range(10)]
  
      edge_label_index = graphs[0].edge_label_index
      dataset = GraphDataset(
          graphs,
          task="link_pred",
          edge_message_ratio=0.6,
          edge_train_mode="all",
      )
      datasets = {}
      (
          datasets["train"],
          datasets["val"],
          datasets["test"],
      ) = dataset.split(
          transductive=False, split_ratio=[0.85, 0.05, 0.1]
      )
      edge_label_index_split = datasets["train"][0].edge_label_index[
          :, 0 : edge_label_index.shape[1]
      ]
  
      self.assertTrue(
          torch.equal(edge_label_index, edge_label_index_split)
      )
  
      # transductive split with node task (pytorch geometric dataset)
      pyg_dataset = Planetoid("./cora", "Cora")
      ds = pyg_to_dicts(pyg_dataset, task="cora")
      graphs = [Graph(**item) for item in ds]
      split_ratio = [0.3, 0.3, 0.4]
      node_size_list = [0 for i in range(len(split_ratio))]
      for graph in graphs:
          custom_splits = [[] for i in range(len(split_ratio))]
          split_offset = 0
          num_nodes = graph.num_nodes
          shuffled_node_indices = torch.randperm(graph.num_nodes)
          for i, split_ratio_i in enumerate(split_ratio):
              if i != len(split_ratio) - 1:
                  num_split_i = int(split_ratio_i * num_nodes)
                  nodes_split_i = shuffled_node_indices[
                      split_offset : split_offset + num_split_i
                  ]
                  split_offset += num_split_i
              else:
                  nodes_split_i = shuffled_node_indices[split_offset:]
  
              custom_splits[i] = nodes_split_i
              node_size_list[i] += len(nodes_split_i)
          graph.custom = {"general_splits": custom_splits}
  
      node_feature = graphs[0].node_feature
      edge_index = graphs[0].edge_index
      directed = graphs[0].directed
  
      graph_train = Graph(
          node_feature=node_feature,
          edge_index=edge_index,
          directed=directed,
          node_label_index=graphs[0].custom["general_splits"][0],
      )
  
      graph_val = Graph(
          node_feature=node_feature,
          edge_index=edge_index,
          directed=directed,
          node_label_index=graphs[0].custom["general_splits"][1],
      )
  
      graph_test = Graph(
          node_feature=node_feature,
          edge_index=edge_index,
          directed=directed,
          node_label_index=graphs[0].custom["general_splits"][2],
      )
  
      train_dataset = GraphDataset([graph_train], task="node")
      val_dataset = GraphDataset([graph_val], task="node")
      test_dataset = GraphDataset([graph_test], task="node")
  
      self.assertEqual(
          len(train_dataset[0].node_label_index), node_size_list[0]
      )
      self.assertEqual(
          len(val_dataset[0].node_label_index), node_size_list[1]
      )
      self.assertEqual(
          len(test_dataset[0].node_label_index), node_size_list[2]
      )
  
      # transductive split with edge task
      pyg_dataset = Planetoid("./cora", "Cora")
      graphs_g = GraphDataset.pyg_to_graphs(pyg_dataset)
      ds = pyg_to_dicts(pyg_dataset, task="cora")
      graphs = [Graph(**item) for item in ds]
      split_ratio = [0.3, 0.3, 0.4]
      edge_size_list = [0 for i in range(len(split_ratio))]
      for i, graph in enumerate(graphs):
          custom_splits = [[] for i in range(len(split_ratio))]
          split_offset = 0
          edges = list(graphs_g[i].G.edges)
          num_edges = graph.num_edges
          random.shuffle(edges)
          for i, split_ratio_i in enumerate(split_ratio):
              if i != len(split_ratio) - 1:
                  num_split_i = int(split_ratio_i * num_edges)
                  edges_split_i = edges[
                      split_offset : split_offset + num_split_i
                  ]
                  split_offset += num_split_i
              else:
                  edges_split_i = edges[split_offset:]
  
              custom_splits[i] = edges_split_i
              edge_size_list[i] += len(edges_split_i)
          graph.custom = {"general_splits": custom_splits}
  
      node_feature = graphs[0].node_feature
      edge_index = graphs[0].edge_index
      directed = graphs[0].directed
  
      train_index = torch.tensor(
          graphs[0].custom["general_splits"][0]
      ).permute(1, 0)
      train_index = torch.cat((train_index, train_index), dim=1)
      val_index = torch.tensor(
          graphs[0].custom["general_splits"][1]
      ).permute(1, 0)
      val_index = torch.cat((val_index, val_index), dim=1)
      test_index = torch.tensor(
          graphs[0].custom["general_splits"][2]
      ).permute(1, 0)
      test_index = torch.cat((test_index, test_index), dim=1)
  
      graph_train = Graph(
          node_feature=node_feature,
          edge_index=edge_index,
          directed=directed,
          edge_label_index=train_index,
      )
  
      graph_val = Graph(
          node_feature=node_feature,
          edge_index=edge_index,
          directed=directed,
          edge_label_index=val_index,
      )
  
      graph_test = Graph(
          node_feature=node_feature,
          edge_index=edge_index,
          directed=directed,
          edge_label_index=test_index,
      )
  
      train_dataset = GraphDataset([graph_train], task="edge")
      val_dataset = GraphDataset([graph_val], task="edge")
      test_dataset = GraphDataset([graph_test], task="edge")
  
      self.assertEqual(
          train_dataset[0].edge_label_index.shape[1],
          2 * edge_size_list[0],
      )
      self.assertEqual(
          val_dataset[0].edge_label_index.shape[1],
          2 * edge_size_list[1],
      )
      self.assertEqual(
          test_dataset[0].edge_label_index.shape[1],
          2 * edge_size_list[2],
      )
  
      # inductive split with graph task
      pyg_dataset = TUDataset("./enzymes", "ENZYMES")
      ds = pyg_to_dicts(pyg_dataset)
      graphs = [Graph(**item) for item in ds]
      num_graphs = len(graphs)
      split_ratio = [0.3, 0.3, 0.4]
      graph_size_list = []
      split_offset = 0
      custom_split_graphs = []
      for i, split_ratio_i in enumerate(split_ratio):
          if i != len(split_ratio) - 1:
              num_split_i = int(split_ratio_i * num_graphs)
              custom_split_graphs.append(
                  graphs[split_offset : split_offset + num_split_i]
              )
              split_offset += num_split_i
              graph_size_list.append(num_split_i)
          else:
              custom_split_graphs.append(graphs[split_offset:])
              graph_size_list.append(len(graphs[split_offset:]))
      dataset = GraphDataset(
          graphs, task="graph", custom_split_graphs=custom_split_graphs
      )
      split_res = dataset.split(transductive=False)
      self.assertEqual(graph_size_list[0], len(split_res[0]))
      self.assertEqual(graph_size_list[1], len(split_res[1]))
      self.assertEqual(graph_size_list[2], len(split_res[2]))
  
imports: import copy
  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 pyg_to_dicts, simple_networkx_graph, simple_networkx_small_graph, simple_networkx_graph_alphabet, simple_networkx_multigraph, generate_dense_hete_dataset, generate_simple_small_hete_graph, gen_graph
label_desc: Use image node_id 5 for calling the TestDatasetTensorBackend obj's underlying member method code with example usage: obj.test_dataset_split_custom() without return types