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
ingraph
by feature at indexsortby
.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 callingtorch.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]]]))