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)`(:math:`(B, K, D)\) 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]]]))