code: def test_hetero_graph_split(self):
      # directed G
      G = generate_dense_hete_graph()
      hete = HeteroGraph(G)
      hete = HeteroGraph(
          node_feature=hete.node_feature,
          node_label=hete.node_label,
          edge_feature=hete.edge_feature,
          edge_label=hete.edge_label,
          edge_index=hete.edge_index,
      )
  
      # node
      hete_node = hete.split()
      for node_type in hete.node_label_index:
          num_nodes = 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(hete_node[0].node_label_index[node_type]),
              node_0,
          )
          self.assertEqual(
              len(hete_node[1].node_label_index[node_type]),
              node_1,
          )
          self.assertEqual(
              len(hete_node[2].node_label_index[node_type]),
              node_2,
          )
  
      # node with specified split type
      node_split_types = ["n1"]
      hete_node = hete.split(split_types=node_split_types)
      for node_type in hete.node_label_index:
          if node_type in node_split_types:
              num_nodes = 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(hete_node[0].node_label_index[node_type]),
                  node_0,
              )
              self.assertEqual(
                  len(hete_node[1].node_label_index[node_type]),
                  node_1,
              )
              self.assertEqual(
                  len(hete_node[2].node_label_index[node_type]),
                  node_2,
              )
          else:
              self.assertEqual(
                  len(hete_node[0].node_label_index[node_type]),
                  len(hete.node_label_index[node_type]),
              )
              self.assertEqual(
                  len(hete_node[1].node_label_index[node_type]),
                  len(hete.node_label_index[node_type]),
              )
              self.assertEqual(
                  len(hete_node[2].node_label_index[node_type]),
                  len(hete.node_label_index[node_type]),
              )
  
      # edge
      hete_edge = hete.split(task="edge")
      for edge_type in hete.edge_label_index:
          num_edges = int(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(
              hete_edge[0].edge_label_index[edge_type].shape[1],
              edge_0,
          )
          self.assertEqual(
              hete_edge[1].edge_label_index[edge_type].shape[1],
              edge_1,
          )
          self.assertEqual(
              hete_edge[2].edge_label_index[edge_type].shape[1],
              edge_2,
          )
  
      # edge with specified split type
      edge_split_types = [("n1", "e1", "n1"), ("n1", "e2", "n2")]
      hete_edge = hete.split(task="edge", split_types=edge_split_types)
      for edge_type in hete.edge_label_index:
          if edge_type in edge_split_types:
              num_edges = int(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(
                  hete_edge[0].edge_label_index[edge_type].shape[1],
                  edge_0,
              )
              self.assertEqual(
                  hete_edge[1].edge_label_index[edge_type].shape[1],
                  edge_1,
              )
              self.assertEqual(
                  hete_edge[2].edge_label_index[edge_type].shape[1],
                  edge_2,
              )
          else:
              self.assertEqual(
                  hete_edge[0].edge_label_index[edge_type].shape[1],
                  hete.edge_label_index[edge_type].shape[1],
              )
              self.assertEqual(
                  hete_edge[1].edge_label_index[edge_type].shape[1],
                  hete.edge_label_index[edge_type].shape[1],
              )
              self.assertEqual(
                  hete_edge[2].edge_label_index[edge_type].shape[1],
                  hete.edge_label_index[edge_type].shape[1],
              )
  
      # link_pred
      hete_link = hete.split(
          task="link_pred", split_ratio=[0.5, 0.3, 0.2]
      )
      for key, val in hete.edge_label_index.items():
          num_edges = val.shape[1]
          edge_0 = int(0.5 * num_edges)
          edge_1 = int(0.3 * num_edges)
          edge_2 = num_edges - edge_0 - edge_1
  
          self.assertEqual(
              hete_link[0].edge_label[key].shape[0], edge_0
          )
          self.assertEqual(
              hete_link[1].edge_label[key].shape[0], edge_1
          )
          self.assertEqual(
              hete_link[2].edge_label[key].shape[0], edge_2
          )
  
      # undirected G
      G = generate_dense_hete_graph(directed=False)
      hete = HeteroGraph(G)
      hete = HeteroGraph(
          node_feature=hete.node_feature,
          node_label=hete.node_label,
          edge_feature=hete.edge_feature,
          edge_label=hete.edge_label,
          edge_index=hete.edge_index,
          directed=False,
      )
  
      # node
      hete_node = hete.split()
      for node_type in hete.node_label_index:
          num_nodes = 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(hete_node[0].node_label_index[node_type]),
              node_0,
          )
          self.assertEqual(
              len(hete_node[1].node_label_index[node_type]),
              node_1,
          )
          self.assertEqual(
              len(hete_node[2].node_label_index[node_type]),
              node_2,
          )
  
      # node with specified split type
      node_split_types = ["n1"]
      hete_node = hete.split(split_types=node_split_types)
      for node_type in hete.node_label_index:
          if node_type in node_split_types:
              num_nodes = 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(hete_node[0].node_label_index[node_type]),
                  node_0,
              )
              self.assertEqual(
                  len(hete_node[1].node_label_index[node_type]),
                  node_1,
              )
              self.assertEqual(
                  len(hete_node[2].node_label_index[node_type]),
                  node_2,
              )
          else:
              self.assertEqual(
                  len(hete_node[0].node_label_index[node_type]),
                  len(hete.node_label_index[node_type]),
              )
              self.assertEqual(
                  len(hete_node[1].node_label_index[node_type]),
                  len(hete.node_label_index[node_type]),
              )
              self.assertEqual(
                  len(hete_node[2].node_label_index[node_type]),
                  len(hete.node_label_index[node_type]),
              )
  
      # edge
      hete_edge = hete.split(task="edge")
      for edge_type in hete.edge_label_index:
          num_edges = int(hete.num_edges(edge_type))
          edge_0 = int(num_edges * 0.8)
          edge_1 = int(num_edges * 0.1)
          edge_2 = num_edges - edge_0 - edge_1
          self.assertEqual(
              hete_edge[0].edge_label_index[edge_type].shape[1],
              edge_0,
          )
          self.assertEqual(
              hete_edge[1].edge_label_index[edge_type].shape[1],
              edge_1,
          )
          self.assertEqual(
              hete_edge[2].edge_label_index[edge_type].shape[1],
              edge_2,
          )
  
      # edge with specified split type
      edge_split_types = [("n1", "e1", "n1"), ("n1", "e2", "n2")]
      hete_edge = hete.split(task="edge", split_types=edge_split_types)
      for edge_type in hete.edge_label_index:
          if edge_type in edge_split_types:
              num_edges = int(hete.num_edges(edge_type))
              edge_0 = int(num_edges * 0.8)
              edge_1 = int(num_edges * 0.1)
              edge_2 = num_edges - edge_0 - edge_1
  
              self.assertEqual(
                  hete_edge[0].edge_label_index[edge_type].shape[1],
                  edge_0,
              )
              self.assertEqual(
                  hete_edge[1].edge_label_index[edge_type].shape[1],
                  edge_1,
              )
              self.assertEqual(
                  hete_edge[2].edge_label_index[edge_type].shape[1],
                  edge_2,
              )
          else:
              self.assertEqual(
                  hete_edge[0].edge_label_index[edge_type].shape[1],
                  hete.edge_label_index[edge_type].shape[1],
              )
              self.assertEqual(
                  hete_edge[1].edge_label_index[edge_type].shape[1],
                  hete.edge_label_index[edge_type].shape[1],
              )
              self.assertEqual(
                  hete_edge[2].edge_label_index[edge_type].shape[1],
                  hete.edge_label_index[edge_type].shape[1],
              )
  
      # link_pred
      hete_link = hete.split(
          task="link_pred", split_ratio=[0.5, 0.3, 0.2]
      )
      for key, val in hete.edge_label_index.items():
          num_edges = int(val.shape[1] / 2)
          edge_0 = 2 * int(0.5 * num_edges)
          edge_1 = 2 * int(0.3 * num_edges)
          edge_2 = 2 * (
              num_edges - int(0.5 * num_edges) - int(0.3 * num_edges)
          )
  
          self.assertEqual(
              hete_link[0].edge_label[key].shape[0], edge_0
          )
          self.assertEqual(
              hete_link[1].edge_label[key].shape[0], edge_1
          )
          self.assertEqual(
              hete_link[2].edge_label[key].shape[0], edge_2
          )
  
file_path: deepsnap/tests/test_hetero_graph_tensor_backend.py
imports: import math
  import unittest
  from torch.utils.data import DataLoader
  from tests.utils import generate_simple_hete_graph, generate_dense_hete_graph, generate_dense_hete_multigraph
  from deepsnap.batch import Batch
  from deepsnap.hetero_graph import HeteroGraph
label_desc: Use image node_id 4 for calling the TestHeteroGraphTensorBackend obj's underlying member method code with example usage: obj.test_hetero_graph_split() without return types