update nnlib

Author: CarloLucibello

GNNlib/Project.toml

@@ -7,7 +7,6 @@ version = "1.0.0"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c"
-GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
@@ -32,7 +31,6 @@ CUDA = "5"
 ChainRulesCore = "1.24"
 DataStructures = "0.18"
 GNNGraphs = "1.4"
-GPUArraysCore = "0.1"
 LinearAlgebra = "1"
 MLUtils = "0.4"
 NNlib = "0.9"

GNNlib/src/layers/pool.jl

@@ -29,8 +29,8 @@ topk_index(y::Adjoint, k::Int) = topk_index(y', k)
 function set2set_pool(l, g::GNNGraph, x::AbstractMatrix)
     n_in = size(x, 1)    
     qstar = zeros_like(x, (2*n_in, g.num_graphs))
-    h = zeros_like(l.Wh, size(l.Wh, 2))
-    c = zeros_like(l.Wh, size(l.Wh, 2))
+    h = zeros_like(l.lstm.Wh, size(l.lstm.Wh, 2))
+    c = zeros_like(l.lstm.Wh, size(l.lstm.Wh, 2))
     state = (h, c)
     for t in 1:l.num_iters
         q, state = l.lstm(qstar, state)                     # [n_in, n_graphs]

GraphNeuralNetworks/src/layers/pool.jl

@@ -155,9 +155,6 @@ function Set2Set(n_in::Int, n_iters::Int, n_layers::Int = 1)
     return Set2Set(lstm, n_iters)
 end
 
-function (l::Set2Set)(g, x)
-    m = (; l.lstm, l.num_iters, Wh = l.lstm.Wh)
-    return GNNlib.set2set_pool(m, g, x)
-end
+(l::Set2Set)(g, x) = GNNlib.set2set_pool(l, g, x)
 
 (l::Set2Set)(g::GNNGraph) = GNNGraph(g, gdata = l(g, node_features(g)))