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
Create two
DGLGraph
objects and initialize their edge features.>>> 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.