[GraphBolt] Remove `unique_and_compact_node_pairs`. (#6859)

Co-authored-by: Ubuntu <ubuntu@ip-172-31-0-133.us-west-2.compute.internal>

python/dgl/graphbolt/__init__.py

@@ -22,7 +22,6 @@ from .internal import (
     compact_csc_format,
     unique_and_compact,
     unique_and_compact_csc_formats,
-    unique_and_compact_node_pairs,
 )
 from .utils import add_reverse_edges, exclude_seed_edges
 

python/dgl/graphbolt/internal/sample_utils.py

@@ -61,123 +61,6 @@ def unique_and_compact(
         return unique_and_compact_per_type(nodes)
 
 
-def unique_and_compact_node_pairs(
-    node_pairs: Union[
-        Tuple[torch.Tensor, torch.Tensor],
-        Dict[str, Tuple[torch.Tensor, torch.Tensor]],
-    ],
-    unique_dst_nodes: Union[
-        torch.Tensor,
-        Dict[str, torch.Tensor],
-    ] = None,
-):
-    """
-    Compact node pairs and return unique nodes (per type).
-
-    Parameters
-    ----------
-    node_pairs : Union[Tuple[torch.Tensor, torch.Tensor],
-                    Dict(str, Tuple[torch.Tensor, torch.Tensor])]
-        Node pairs representing source-destination edges.
-        - If `node_pairs` is a tuple: It means the graph is homogeneous.
-        Also, it should be in the format ('u', 'v') representing source
-        and destination pairs. And IDs inside are homogeneous ids.
-        - If `node_pairs` is a dictionary: The keys should be edge type and
-        the values should be corresponding node pairs. And IDs inside are
-        heterogeneous ids.
-    unique_dst_nodes: torch.Tensor or Dict[str, torch.Tensor]
-        Unique nodes of all destination nodes in the node pairs.
-        - If `unique_dst_nodes` is a tensor: It means the graph is homogeneous.
-        - If `node_pairs` is a dictionary: The keys are node type and the
-        values are corresponding nodes. And IDs inside are heterogeneous ids.
-
-    Returns
-    -------
-    Tuple[node_pairs, unique_nodes]
-        The compacted node pairs, where node IDs are replaced with mapped node
-        IDs, and the unique nodes (per type).
-        "Compacted node pairs" indicates that the node IDs in the input node
-        pairs are replaced with mapped node IDs, where each type of node is
-        mapped to a contiguous space of IDs ranging from 0 to N.
-
-    Examples
-    --------
-    >>> import dgl.graphbolt as gb
-    >>> N1 = torch.LongTensor([1, 2, 2])
-    >>> N2 = torch.LongTensor([5, 6, 5])
-    >>> node_pairs = {"n1:e1:n2": (N1, N2),
-    ...     "n2:e2:n1": (N2, N1)}
-    >>> unique_nodes, compacted_node_pairs = gb.unique_and_compact_node_pairs(
-    ...     node_pairs
-    ... )
-    >>> print(unique_nodes)
-    {'n1': tensor([1, 2]), 'n2': tensor([5, 6])}
-    >>> print(compacted_node_pairs)
-    {"n1:e1:n2": (tensor([0, 1, 1]), tensor([0, 1, 0])),
-    "n2:e2:n1": (tensor([0, 1, 0]), tensor([0, 1, 1]))}
-    """
-    is_homogeneous = not isinstance(node_pairs, dict)
-    if is_homogeneous:
-        node_pairs = {"_N:_E:_N": node_pairs}
-        if unique_dst_nodes is not None:
-            assert isinstance(
-                unique_dst_nodes, torch.Tensor
-            ), "Edge type not supported in homogeneous graph."
-            unique_dst_nodes = {"_N": unique_dst_nodes}
-
-    # Collect all source and destination nodes for each node type.
-    src_nodes = defaultdict(list)
-    dst_nodes = defaultdict(list)
-    device = None
-    for etype, (src_node, dst_node) in node_pairs.items():
-        if device is None:
-            device = src_node.device
-        src_type, _, dst_type = etype_str_to_tuple(etype)
-        src_nodes[src_type].append(src_node)
-        dst_nodes[dst_type].append(dst_node)
-    src_nodes = {ntype: torch.cat(nodes) for ntype, nodes in src_nodes.items()}
-    dst_nodes = {ntype: torch.cat(nodes) for ntype, nodes in dst_nodes.items()}
-    # Compute unique destination nodes if not provided.
-    if unique_dst_nodes is None:
-        unique_dst_nodes = {
-            ntype: torch.unique(nodes) for ntype, nodes in dst_nodes.items()
-        }
-
-    ntypes = set(dst_nodes.keys()) | set(src_nodes.keys())
-    unique_nodes = {}
-    compacted_src = {}
-    compacted_dst = {}
-    dtype = list(src_nodes.values())[0].dtype
-    default_tensor = torch.tensor([], dtype=dtype, device=device)
-    for ntype in ntypes:
-        src = src_nodes.get(ntype, default_tensor)
-        unique_dst = unique_dst_nodes.get(ntype, default_tensor)
-        dst = dst_nodes.get(ntype, default_tensor)
-        (
-            unique_nodes[ntype],
-            compacted_src[ntype],
-            compacted_dst[ntype],
-        ) = torch.ops.graphbolt.unique_and_compact(src, dst, unique_dst)
-
-    compacted_node_pairs = {}
-    # Map back with the same order.
-    for etype, pair in node_pairs.items():
-        num_elem = pair[0].size(0)
-        src_type, _, dst_type = etype_str_to_tuple(etype)
-        src = compacted_src[src_type][:num_elem]
-        dst = compacted_dst[dst_type][:num_elem]
-        compacted_node_pairs[etype] = (src, dst)
-        compacted_src[src_type] = compacted_src[src_type][num_elem:]
-        compacted_dst[dst_type] = compacted_dst[dst_type][num_elem:]
-
-    # Return singleton for a homogeneous graph.
-    if is_homogeneous:
-        compacted_node_pairs = list(compacted_node_pairs.values())[0]
-        unique_nodes = list(unique_nodes.values())[0]
-
-    return unique_nodes, compacted_node_pairs
-
-
 def unique_and_compact_csc_formats(
     csc_formats: Union[
         Tuple[torch.Tensor, torch.Tensor],

tests/python/pytorch/graphbolt/test_graphbolt_utils.py

@@ -160,109 +160,6 @@ def test_unique_and_compact_homo():
         assert torch.equal(expected_node, node)
 
 
-def test_unique_and_compact_node_pairs_hetero():
-    node_pairs = {
-        "n1:e1:n2": (
-            torch.tensor([1, 3, 4, 6, 2, 7, 9, 4, 2, 6], device=F.ctx()),
-            torch.tensor([2, 2, 2, 4, 1, 1, 1, 3, 3, 3], device=F.ctx()),
-        ),
-        "n1:e2:n3": (
-            torch.tensor([5, 2, 6, 4, 7, 2, 8, 1, 3, 0], device=F.ctx()),
-            torch.tensor([1, 3, 3, 3, 2, 2, 2, 7, 7, 7], device=F.ctx()),
-        ),
-        "n2:e3:n3": (
-            torch.tensor([2, 5, 4, 1, 4, 3, 6, 0], device=F.ctx()),
-            torch.tensor([1, 1, 3, 3, 2, 2, 7, 7], device=F.ctx()),
-        ),
-    }
-
-    expected_unique_nodes = {
-        "n1": torch.tensor([1, 3, 4, 6, 2, 7, 9, 5, 8, 0], device=F.ctx()),
-        "n2": torch.tensor([1, 2, 3, 4, 5, 6, 0], device=F.ctx()),
-        "n3": torch.tensor([1, 2, 3, 7], device=F.ctx()),
-    }
-    if expected_unique_nodes["n1"].is_cuda:
-        expected_reverse_id = {
-            "n1": expected_unique_nodes["n1"].sort()[1],
-            "n2": torch.tensor([0, 1, 2, 3, 6, 4, 5], device=F.ctx()),
-            "n3": expected_unique_nodes["n3"].sort()[1],
-        }
-        expected_unique_nodes = {
-            "n1": expected_unique_nodes["n1"].sort()[0],
-            "n2": torch.tensor([1, 2, 3, 4, 0, 5, 6], device=F.ctx()),
-            "n3": expected_unique_nodes["n3"].sort()[0],
-        }
-    else:
-        expected_reverse_id = {
-            k: torch.arange(0, v.shape[0], device=F.ctx())
-            for k, v in expected_unique_nodes.items()
-        }
-    expected_node_pairs = {
-        "n1:e1:n2": (
-            torch.tensor([0, 1, 2, 3, 4, 5, 6, 2, 4, 3], device=F.ctx()),
-            torch.tensor([1, 1, 1, 3, 0, 0, 0, 2, 2, 2], device=F.ctx()),
-        ),
-        "n1:e2:n3": (
-            torch.tensor([7, 4, 3, 2, 5, 4, 8, 0, 1, 9], device=F.ctx()),
-            torch.tensor([0, 2, 2, 2, 1, 1, 1, 3, 3, 3], device=F.ctx()),
-        ),
-        "n2:e3:n3": (
-            torch.tensor([1, 4, 3, 0, 3, 2, 5, 6], device=F.ctx()),
-            torch.tensor([0, 0, 2, 2, 1, 1, 3, 3], device=F.ctx()),
-        ),
-    }
-
-    unique_nodes, compacted_node_pairs = gb.unique_and_compact_node_pairs(
-        node_pairs
-    )
-    for ntype, nodes in unique_nodes.items():
-        expected_nodes = expected_unique_nodes[ntype]
-        assert torch.equal(nodes, expected_nodes)
-    for etype, pair in compacted_node_pairs.items():
-        u, v = pair
-        ntype1, _, ntype2 = etype.split(":")
-        u = expected_reverse_id[ntype1][u]
-        v = expected_reverse_id[ntype2][v]
-        expected_u, expected_v = expected_node_pairs[etype]
-        assert torch.equal(u, expected_u)
-        assert torch.equal(v, expected_v)
-
-
-def test_unique_and_compact_node_pairs_homo():
-    dst_nodes = torch.tensor([1, 1, 3, 3, 5, 5, 2, 6, 6, 6, 6], device=F.ctx())
-    src_nodes = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6], device=F.ctx())
-    node_pairs = (src_nodes, dst_nodes)
-
-    expected_unique_nodes = torch.tensor([1, 2, 3, 5, 6, 4], device=F.ctx())
-    expected_dst_nodes = torch.tensor(
-        [0, 0, 2, 2, 3, 3, 1, 4, 4, 4, 4], device=F.ctx()
-    )
-    expected_src_ndoes = torch.tensor(
-        [1, 2, 0, 5, 3, 1, 3, 0, 5, 5, 4], device=F.ctx()
-    )
-    unique_nodes, compacted_node_pairs = gb.unique_and_compact_node_pairs(
-        node_pairs
-    )
-    assert torch.equal(unique_nodes, expected_unique_nodes)
-
-    u, v = compacted_node_pairs
-    assert torch.equal(u, expected_src_ndoes)
-    assert torch.equal(v, expected_dst_nodes)
-    assert torch.equal(
-        unique_nodes[:5], torch.tensor([1, 2, 3, 5, 6], device=F.ctx())
-    )
-
-
-def test_incomplete_unique_dst_nodes_():
-    node_pairs = (
-        torch.arange(0, 50, device=F.ctx()),
-        torch.arange(100, 150, device=F.ctx()),
-    )
-    unique_dst_nodes = torch.arange(150, 200, device=F.ctx())
-    with pytest.raises(IndexError):
-        gb.unique_and_compact_node_pairs(node_pairs, unique_dst_nodes)
-
-
 def test_unique_and_compact_csc_formats_hetero():
     dst_nodes = {
         "n2": torch.tensor([2, 4, 1, 3]),