MCPcopy Index your code
hub / github.com/dmlc/dgl / set_batch_num_edges

Method set_batch_num_edges

python/dgl/heterograph.py:1678–1763  ·  view source on GitHub ↗

Manually set the number of edges for each graph in the batch with the specified edge type. Parameters ---------- val : Tensor or Mapping[str, Tensor] The dictionary storing number of edges for each graph in the batch for all edge types. If the

(self, val)

Source from the content-addressed store, hash-verified

1676 return self._batch_num_edges[etype]
1677
1678 def set_batch_num_edges(self, val):
1679 """Manually set the number of edges for each graph in the batch with the specified edge
1680 type.
1681
1682 Parameters
1683 ----------
1684 val : Tensor or Mapping[str, Tensor]
1685 The dictionary storing number of edges for each graph in the batch for all edge types.
1686 If the graph has only one edge type, ``val`` can also be a single array indicating the
1687 number of edges per graph in the batch.
1688
1689 Notes
1690 -----
1691 This API is always used together with ``set_batch_num_nodes`` to specify batching
1692 information of a graph, it also do not check the correspondance between the graph structure
1693 and batching information and user must guarantee there will be no cross-graph edges in the
1694 batch.
1695
1696 Examples
1697 --------
1698
1699 The following example uses PyTorch backend.
1700
1701 >>> import dgl
1702 >>> import torch
1703
1704 Create a homogeneous graph.
1705
1706 >>> g = dgl.graph(([0, 1, 2, 3, 4, 5], [1, 2, 0, 4, 5, 3]))
1707
1708 Manually set batch information
1709
1710 >>> g.set_batch_num_nodes(torch.tensor([3, 3]))
1711 >>> g.set_batch_num_edges(torch.tensor([3, 3]))
1712
1713 Unbatch the graph.
1714
1715 >>> dgl.unbatch(g)
1716 [Graph(num_nodes=3, num_edges=3,
1717 ndata_schemes={}
1718 edata_schemes={}), Graph(num_nodes=3, num_edges=3,
1719 ndata_schemes={}
1720 edata_schemes={})]
1721
1722 Create a heterogeneous graph.
1723
1724 >>> hg = dgl.heterograph({
1725 ... ('user', 'plays', 'game') : ([0, 1, 2, 3, 4, 5], [0, 1, 1, 3, 3, 2]),
1726 ... ('developer', 'develops', 'game') : ([0, 1, 2, 3], [1, 0, 3, 2])})
1727
1728 Manually set batch information.
1729
1730 >>> hg.set_batch_num_nodes({
1731 ... 'user': torch.tensor([3, 3]),
1732 ... 'game': torch.tensor([2, 2]),
1733 ... 'developer': torch.tensor([2, 2])})
1734 >>> hg.set_batch_num_edges(
1735 ... {('user', 'plays', 'game'): torch.tensor([3, 3]),

Callers 5

batchFunction · 0.95
knn_graphFunction · 0.80
segmented_knn_graphFunction · 0.80
test_set_batch_infoFunction · 0.80
test_remove_selfloopFunction · 0.80

Calls 1

DGLErrorClass · 0.85

Tested by 2

test_set_batch_infoFunction · 0.64
test_remove_selfloopFunction · 0.64