[GraphBolt] use str etype in data block (#6231)

python/dgl/graphbolt/data_block.py

@@ -35,12 +35,9 @@ class DataBlock:
     The keys are tuples in the format '(edge_type, feature_name)', and the
     values represent the corresponding features. In the case of a homogeneous
     graph where no edge types exist, 'edge_type' should be set to None.
-    Note 'edge_type' are of format 'str:str:str'.
     """
 
-    input_nodes: Union[
-        torch.Tensor, Dict[Tuple[str, str, str], torch.Tensor]
-    ] = None
+    input_nodes: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
     """A representation of input nodes in the outermost layer. Conatins all nodes
        in the 'sampled_subgraphs'.
     - If `input_nodes` is a tensor: It indicates the graph is homogeneous.

python/dgl/graphbolt/link_prediction_block.py

@@ -16,7 +16,7 @@ class LinkPredictionBlock(DataBlock):
 
     node_pair: Union[
         Tuple[torch.Tensor, torch.Tensor],
-        Dict[Tuple[str, str, str], Tuple[torch.Tensor, torch.Tensor]],
+        Dict[str, Tuple[torch.Tensor, torch.Tensor]],
     ] = None
     """
     Representation of seed node pairs utilized in link prediction tasks.
@@ -27,7 +27,7 @@ class LinkPredictionBlock(DataBlock):
       type.
     """
 
-    label: Union[torch.Tensor, Dict[Tuple[str, str, str], torch.Tensor]] = None
+    label: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
     """
     Labels associated with the link prediction task.
     - If `label` is a tensor: It indicates a homogeneous graph. The value are
@@ -36,9 +36,7 @@ class LinkPredictionBlock(DataBlock):
       should correspond to given 'node_pair'.
     """
 
-    negative_head: Union[
-        torch.Tensor, Dict[Tuple[str, str, str], torch.Tensor]
-    ] = None
+    negative_head: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
     """
     Representation of negative samples for the head nodes in the link
     prediction task.
@@ -48,9 +46,7 @@ class LinkPredictionBlock(DataBlock):
       given type.
     """
 
-    negative_tail: Union[
-        torch.Tensor, Dict[Tuple[str, str, str], torch.Tensor]
-    ] = None
+    negative_tail: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
     """
     Representation of negative samples for the tail nodes in the link
     prediction task.
@@ -62,24 +58,20 @@ class LinkPredictionBlock(DataBlock):
 
     compacted_node_pair: Union[
         Tuple[torch.Tensor, torch.Tensor],
-        Dict[Tuple[str, str, str], Tuple[torch.Tensor, torch.Tensor]],
+        Dict[str, Tuple[torch.Tensor, torch.Tensor]],
     ] = None
     """
     Representation of compacted node pairs corresponding to 'node_pair', where
     all node ids inside are compacted.
     """
 
-    compacted_negative_head: Union[
-        torch.Tensor, Dict[Tuple[str, str, str], torch.Tensor]
-    ] = None
+    compacted_negative_head: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
     """
     Representation of compacted nodes corresponding to 'negative_head', where
     all node ids inside are compacted.
     """
 
-    compacted_negative_tail: Union[
-        torch.Tensor, Dict[Tuple[str, str, str], torch.Tensor]
-    ] = None
+    compacted_negative_tail: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
     """
     Representation of compacted nodes corresponding to 'negative_tail', where
     all node ids inside are compacted.