dgl.topk_edges¶

dgl.topk_edges(graph, feat, k, *, descending=True, sortby=None, etype=None)[source]¶

Return a graph-level representation by a graph-wise top-k on edge features feat in graph by feature at index sortby.

If descending is set to False, return the k smallest elements instead.

If sortby is set to None, the function would perform top-k on all dimensions independently, equivalent to calling torch.topk(graph.edata[feat], dim=0).

Parameters

graph (DGLGraph) – The graph.
feat (str) – The feature field.
k (int) – The k in “top-k”
descending (bool) – Controls whether to return the largest or smallest elements.
sortby (int, optional) – Sort according to which feature. If is None, all features are sorted independently.
etype (str, typle of str, optional) – Edge type. Can be omitted if there is only one edge type in the graph.

Returns

sorted_feat (Tensor) – A tensor with shape \((B, K, D)\), where \(B\) is the batch size of the input graph.
sorted_idx (Tensor) – A tensor with shape \((B, K)\) if sortby is set to None), where \(B\) is the batch size of the input graph, \(D\) is the feature size.

Notes

If an example has \(n\) nodes and \(n<k\), the sorted_feat tensor will pad the \(n+1\) to \(k\) th rows with zero;

Examples

>>> import dgl
>>> import torch as th

Create two DGLGraph objects and initialize their edge features.

>>> g1 = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 0]))         # Graph 1
>>> 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.graph(([0, 1, 2, 3, 4], [1, 2, 3, 4, 0]))   # Graph 2
>>> 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]])

Top-k over edge attribute h in a batched graph.

>>> bg = dgl.batch([g1, g2], edata=['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]]]))

Top-k over edge attribute h along index -1 in a batched graph. (used in SortPooling)

>>> dgl.topk_edges(bg, 'h', 3, sortby=-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]]))

Top-k 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]]]))