-
Notifications
You must be signed in to change notification settings - Fork 45
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add topk features #260
base: master
Are you sure you want to change the base?
Add topk features #260
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.
|
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
the masking would be different for edge feature
With this PR I add the
topk_nodes
andtopk_edges
features (see issue #41).