dgl.topk_edges¶

dgl.
topk_edges
(graph, feat, k, descending=True, idx=None)[source]¶ Return graphwise topk edge features of field
feat
ingraph
ranked by keys at given indexidx
. Ifdescending
is set to False, return the k smallest elements instead.If idx is set to None, the function would return topk value of all indices, which is equivalent to calling
torch.topk(graph.edata[feat], dim=0)
for each example of the input graph.Parameters: Returns: The first tensor returns topk edge features of the given graph with shape \((K, D)\), if the input graph is a BatchedDGLGraph, a tensor with shape \((B, K, D)\) would be returned, where \(B\) is the batch size. The second tensor returns the topk edge indices of the given graph with shape \((K)`(:math:`(K, D)\) if idx is set to None), if the input graph is a BatchedDGLGraph, a tensor with shape \((B, K)`(:math:`(B, K, D)\) if` idx is set to None) would be returned, where \(B\) is the batch size.
Return type: tuple of tensors
Examples
>>> import dgl >>> import torch as th
Create two
DGLGraph
objects and initialize their edge features.>>> g1 = dgl.DGLGraph() # Graph 1 >>> g1.add_nodes(4) >>> g1.add_edges([0, 1, 2, 3], [1, 2, 3, 0]) >>> g1.edata['h'] = th.rand(4, 5) >>> g1.edata['h'] tensor([[0.0297, 0.8307, 0.9140, 0.6702, 0.3346], [0.5901, 0.3030, 0.9280, 0.6893, 0.7997], [0.0880, 0.6515, 0.4451, 0.7507, 0.5297], [0.5171, 0.6379, 0.2695, 0.8954, 0.5197]])
>>> g2 = dgl.DGLGraph() # Graph 2 >>> g2.add_nodes(5) >>> g2.add_edges([0, 1, 2, 3, 4], [1, 2, 3, 4, 0]) >>> g2.edata['h'] = th.rand(5, 5) >>> g2.edata['h'] tensor([[0.3168, 0.3174, 0.5303, 0.0804, 0.3808], [0.1323, 0.2766, 0.4318, 0.6114, 0.1458], [0.1752, 0.9105, 0.5692, 0.8489, 0.0539], [0.1931, 0.4954, 0.3455, 0.3934, 0.0857], [0.5065, 0.5182, 0.5418, 0.1520, 0.3872]])
Topk over edge attribute
h
in a batched graph.>>> bg = dgl.batch([g1, g2], edge_attrs='h') >>> dgl.topk_edges(bg, 'h', 3) (tensor([[[0.5901, 0.8307, 0.9280, 0.8954, 0.7997], [0.5171, 0.6515, 0.9140, 0.7507, 0.5297], [0.0880, 0.6379, 0.4451, 0.6893, 0.5197]],
 [[0.5065, 0.9105, 0.5692, 0.8489, 0.3872],
 [0.3168, 0.5182, 0.5418, 0.6114, 0.3808], [0.1931, 0.4954, 0.5303, 0.3934, 0.1458]]]), tensor([[[1, 0, 1, 3, 1], [3, 2, 0, 2, 2], [2, 3, 2, 1, 3]],
 [[4, 2, 2, 2, 4],
 [0, 4, 4, 1, 0], [3, 3, 0, 3, 1]]]))
Topk over edge attribute
h
along index 1 in a batched graph. (used in SortPooling)>>> dgl.topk_edges(bg, 'h', 3, idx=1) (tensor([[[0.5901, 0.3030, 0.9280, 0.6893, 0.7997], [0.0880, 0.6515, 0.4451, 0.7507, 0.5297], [0.5171, 0.6379, 0.2695, 0.8954, 0.5197]],
 [[0.5065, 0.5182, 0.5418, 0.1520, 0.3872],
 [0.3168, 0.3174, 0.5303, 0.0804, 0.3808], [0.1323, 0.2766, 0.4318, 0.6114, 0.1458]]]), tensor([[1, 2, 3],
[4, 0, 1]]))
Topk over edge attribute
h
in a single graph.>>> dgl.topk_edges(g1, 'h', 3) (tensor([[0.5901, 0.8307, 0.9280, 0.8954, 0.7997], [0.5171, 0.6515, 0.9140, 0.7507, 0.5297], [0.0880, 0.6379, 0.4451, 0.6893, 0.5197]]), tensor([[[1, 0, 1, 3, 1], [3, 2, 0, 2, 2], [2, 3, 2, 1, 3]]]))
Notes
If an example has \(n\) edges and \(n<k\), in the first returned tensor the \(n+1\) to \(k`th rows would be padded with all zero; in the second returned tensor, the behavior of :math:`n+1\) to :math:`k`th elements is not defined.