Open aurorarossi opened 1 year ago
I think it is more useful and general to have a function with signature
topk_nodes(g::GNNGraph, x::AbstractArray, k::Int; rev::Bool = true, sortby::Union{Nothing, Int} = nothing)
instead of passing the tensor name.
Also, the docstring needs much more explanation of the purpose of this function, the various arguments, and the returned objects.
Following DGL, I think it is useful to return both the sorted array and permutation.
So if we call (y, partialperm)
the output:
y
should be an array of size (num_feat, k, num_graphs)
partialperm
should be an array of integers of size (k, num_graphs)
if sortby === nothing
and of size (num_feat, k, num_graphs)
if sortby
is an integer.Does it work on gpu as well?
Yes, it works on the GPU.
I think it would be nice to contribute topk_matrix
as a topk
function in MLUtils.jl and then base this PR on top of that.
Then I would have here two functions, topk_nodes
and topk_edges
as in Deep Graph Library (https://docs.dgl.ai/generated/dgl.topk_nodes.html#dgl.topk_nodes)
With this PR I add the
topk_nodes
andtopk_edges
features (see issue https://github.com/CarloLucibello/GraphNeuralNetworks.jl/issues/41).