chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,131 @@
+"""Module for various graph generator functions."""
+
+from . import backend as F, convert, random
+
+__all__ = ["rand_graph", "rand_bipartite"]
+
+
+def rand_graph(num_nodes, num_edges, idtype=F.int64, device=F.cpu()):
+    """Generate a random graph of the given number of nodes/edges and return.
+
+    It uniformly chooses ``num_edges`` from all possible node pairs and form a graph.
+    The random choice is without replacement, which means there will be no multi-edge
+    in the resulting graph.
+
+    To control the randomness, set the random seed via :func:`dgl.seed`.
+
+    Parameters
+    ----------
+    num_nodes : int
+        The number of nodes
+    num_edges : int
+        The number of edges
+    idtype : int32, int64, optional
+        The data type for storing the structure-related graph information
+        such as node and edge IDs. It should be a framework-specific data type object
+        (e.g., torch.int32). By default, DGL uses int64.
+    device : Device context, optional
+        The device of the resulting graph. It should be a framework-specific device
+        object (e.g., torch.device). By default, DGL stores the graph on CPU.
+
+    Returns
+    -------
+    DGLGraph
+        The generated random graph.
+
+    See Also
+    --------
+    rand_bipartite
+
+    Examples
+    --------
+    >>> import dgl
+    >>> dgl.rand_graph(100, 10)
+    Graph(num_nodes=100, num_edges=10,
+          ndata_schemes={}
+          edata_schemes={})
+    """
+    # TODO(minjie): support RNG as one of the arguments.
+    eids = random.choice(num_nodes * num_nodes, num_edges, replace=False)
+    eids = F.zerocopy_to_numpy(eids)
+    rows = F.zerocopy_from_numpy(eids // num_nodes)
+    cols = F.zerocopy_from_numpy(eids % num_nodes)
+    rows = F.copy_to(F.astype(rows, idtype), device)
+    cols = F.copy_to(F.astype(cols, idtype), device)
+    return convert.graph(
+        (rows, cols), num_nodes=num_nodes, idtype=idtype, device=device
+    )
+
+
+def rand_bipartite(
+    utype,
+    etype,
+    vtype,
+    num_src_nodes,
+    num_dst_nodes,
+    num_edges,
+    idtype=F.int64,
+    device=F.cpu(),
+):
+    """Generate a random uni-directional bipartite graph and return.
+
+    It uniformly chooses ``num_edges`` from all possible node pairs and form a graph.
+    The random choice is without replacement, which means there will be no multi-edge
+    in the resulting graph.
+
+    To control the randomness, set the random seed via :func:`dgl.seed`.
+
+    Parameters
+    ----------
+    utype : str, optional
+        The name of the source node type.
+    etype : str, optional
+        The name of the edge type.
+    vtype : str, optional
+        The name of the destination node type.
+    num_src_nodes : int
+        The number of source nodes.
+    num_dst_nodes : int
+        The number of destination nodes.
+    num_edges : int
+        The number of edges
+    idtype : int32, int64, optional
+        The data type for storing the structure-related graph information
+        such as node and edge IDs. It should be a framework-specific data type object
+        (e.g., torch.int32). By default, DGL uses int64.
+    device : Device context, optional
+        The device of the resulting graph. It should be a framework-specific device
+        object (e.g., torch.device). By default, DGL stores the graph on CPU.
+
+    Returns
+    -------
+    DGLGraph
+        The generated random bipartite graph.
+
+    See Also
+    --------
+    rand_graph
+
+    Examples
+    --------
+    >>> import dgl
+    >>> dgl.rand_bipartite('user', 'buys', 'game', 50, 100, 10)
+    Graph(num_nodes={'game': 100, 'user': 50},
+          num_edges={('user', 'buys', 'game'): 10},
+          metagraph=[('user', 'game', 'buys')])
+    """
+    # TODO(minjie): support RNG as one of the arguments.
+    eids = random.choice(
+        num_src_nodes * num_dst_nodes, num_edges, replace=False
+    )
+    eids = F.zerocopy_to_numpy(eids)
+    rows = F.zerocopy_from_numpy(eids // num_dst_nodes)
+    cols = F.zerocopy_from_numpy(eids % num_dst_nodes)
+    rows = F.copy_to(F.astype(rows, idtype), device)
+    cols = F.copy_to(F.astype(cols, idtype), device)
+    return convert.heterograph(
+        {(utype, etype, vtype): (rows, cols)},
+        {utype: num_src_nodes, vtype: num_dst_nodes},
+        idtype=idtype,
+        device=device,
+    )