Conversation
|
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.
|
aurorarossi
force-pushed
the
add_topk_GPU
branch
from
b02dcaa to
05fca7c
Compare
|
Does it work on gpu as well? |
|
Yes, it works on the GPU. |
|
I think it would be nice to contribute Then I would have here two functions, |
| - `feat`: a feature array of size `(number_features, g.num_nodes)` or `(number_features, g.num_edges)` of the graph `g`. | ||
| - `k`: the number of top features to return. | ||
| - `rev`: if `true`, sort in descending order otherwise returns the `k` smallest elements. | ||
| - `sortby`: the index of the feature to sort by. If `nothing`, every row independently. |
There was a problem hiding this comment.
this sentence is not clear
| function _topk_matrix(matrix::AbstractArray, k::Int; rev::Bool = true, sortby::Union{Nothing, Int} = nothing) | ||
| if sortby === nothing | ||
| sorted_matrix = sort(matrix, dims = 2; rev)[:, 1:k] | ||
| vector_indices = map(x -> sortperm(x; rev), eachrow(matrix)) |
There was a problem hiding this comment.
instead of sorting the whole matrix, it would be more efficient to use partialsortperm. I'm not sure is supported by CUDA.jl though
| if g.num_graphs == 1 | ||
| return _topk_matrix(feat, k; rev, sortby) | ||
| else | ||
| matrices = [feat[:, g.graph_indicator .== i] for i in 1:(g.num_graphs)] |
There was a problem hiding this comment.
the masking would be different for edge feature
2 participants
With this PR I add the
topk_nodesandtopk_edgesfeatures (see issue #41).