renmae message to comput_message

Author: CarloLucibello

src/layers/conv.jl

@@ -45,7 +45,7 @@ function (l::GCNConv)(g::GNNGraph, x::AbstractMatrix{T}) where T
     l.σ.(l.weight * x * Ã .+ l.bias)
 end
 
-message(l::GCNConv, xi, xj, eij) = xj
+compute_message(l::GCNConv, xi, xj, eij) = xj
 update_node(l::GCNConv, m, x) = m
 
 function (l::GCNConv)(g::GNNGraph, x::CuMatrix{T}) where T
@@ -176,7 +176,7 @@ function GraphConv(ch::Pair{Int,Int}, σ=identity, aggr=+;
     GraphConv(W1, W2, b, σ, aggr)
 end
 
-message(l::GraphConv, x_i, x_j, e_ij) =  x_j
+compute_message(l::GraphConv, x_i, x_j, e_ij) =  x_j
 update_node(l::GraphConv, m, x) = l.σ.(l.weight1 * x .+ l.weight2 * m .+ l.bias)
 
 function (l::GraphConv)(g::GNNGraph, x::AbstractMatrix)
@@ -326,7 +326,7 @@ function GatedGraphConv(out_ch::Int, num_layers::Int;
 end
 
 
-message(l::GatedGraphConv, x_i, x_j, e_ij) = x_j
+compute_message(l::GatedGraphConv, x_i, x_j, e_ij) = x_j
 update_node(l::GatedGraphConv, m, x) = m
 
 # remove after https://github.com/JuliaDiff/ChainRules.jl/pull/521
@@ -381,7 +381,7 @@ end
 
 EdgeConv(nn; aggr=max) = EdgeConv(nn, aggr)
 
-message(l::EdgeConv, x_i, x_j, e_ij) = l.nn(vcat(x_i, x_j .- x_i))
+compute_message(l::EdgeConv, x_i, x_j, e_ij) = l.nn(vcat(x_i, x_j .- x_i))
 
 update_node(l::EdgeConv, m, x) = m
 
@@ -426,7 +426,7 @@ function GINConv(nn; eps=0f0)
     GINConv(nn, eps)
 end
 
-message(l::GINConv, x_i, x_j, e_ij) = x_j 
+compute_message(l::GINConv, x_i, x_j, e_ij) = x_j 
 update_node(l::GINConv, m, x) = l.nn((1 + l.eps) * x + m)
 
 function (l::GINConv)(g::GNNGraph, X::AbstractMatrix)

src/msgpass.jl

@@ -12,7 +12,7 @@ with updated feautres.
 The computational steps are the following:
 
 ```julia
-m = compute_batch_message(l, g, x, e)  # calls `message`
+m = compute_batch_message(l, g, x, e)  # calls `compute_message`
 m̄ = aggregate_neighbors(l, aggr, g, m)
 x′ = update_node(l, m̄, x)
 e′ = update_edge(l, m, e)
@@ -72,7 +72,7 @@ end
 ## Step 1.
 
 """
-    message(l, x_i, x_j, [e_ij])
+    compute_message(l, x_i, x_j, [e_ij])
 
 Message function for the message-passing scheme,
 returning the message from node `j` to node `i` .
@@ -81,7 +81,7 @@ from the neighborhood of `i` will later be aggregated
 in order to update (see [`update_node`](@ref)) the features of node `i`.
 
 The function operates on batches of edges, therefore
-`x_i`, `x_j`, and `e_ij` are tensors whose last dimention
+`x_i`, `x_j`, and `e_ij` are tensors whose last dimension
 is the batch size. 
 
 By default, the function returns `x_j`.
@@ -98,8 +98,8 @@ See also [`update_node`](@ref) and [`propagate`](@ref).
 """
 function message end 
 
-@inline message(l, x_i, x_j, e_ij) = message(l, x_i, x_j)
-@inline message(l, x_i, x_j) = x_j
+@inline compute_message(l, x_i, x_j, e_ij) = compute_message(l, x_i, x_j)
+@inline compute_message(l, x_i, x_j) = x_j
 
 _gather(x, i) = NNlib.gather(x, i)
 _gather(x::Nothing, i) = nothing
@@ -108,7 +108,7 @@ function compute_batch_message(l, g, x, e)
     s, t = edge_index(g)
     xi = _gather(x, t)
     xj = _gather(x, s)
-    m = message(l, xi, xj, e)
+    m = compute_message(l, xi, xj, e)
     return m
 end
 

test/cuda/msgpass.jl

@@ -17,7 +17,7 @@
     @functor NewCudaLayer
 
     (l::NewCudaLayer)(g, X) = GraphNeuralNetworks.propagate(l, g, +, X)
-    GraphNeuralNetworks.message(n::NewCudaLayer, x_i, x_j, e_ij) = n.weight * x_j
+    GraphNeuralNetworks.compute_message(n::NewCudaLayer, x_i, x_j, e_ij) = n.weight * x_j
     GraphNeuralNetworks.update_node(::NewCudaLayer, m, x) = m
 
     X = rand(T, in_channel, N) |> gpu

test/msgpass.jl

@@ -51,7 +51,7 @@
     @testset "custom message and neighbor aggregation" begin
         struct NewLayer3{G} end
 
-        GraphNeuralNetworks.message(l::NewLayer3{GRAPH_T}, xi, xj, e) = ones(T, out_channel, size(e,2))
+        GraphNeuralNetworks.compute_message(l::NewLayer3{GRAPH_T}, xi, xj, e) = ones(T, out_channel, size(e,2))
         (l::NewLayer3{GRAPH_T})(g) = GraphNeuralNetworks.propagate(l, g, +)
 
 
@@ -70,7 +70,7 @@
         struct NewLayer4{G} end
 
         GraphNeuralNetworks.update_edge(l::NewLayer4{GRAPH_T}, m, e) = m
-        GraphNeuralNetworks.message(l::NewLayer4{GRAPH_T}, xi, xj, e) = ones(T, out_channel, size(e,2))
+        GraphNeuralNetworks.compute_message(l::NewLayer4{GRAPH_T}, xi, xj, e) = ones(T, out_channel, size(e,2))
         (l::NewLayer4{GRAPH_T})(g) = GraphNeuralNetworks.propagate(l, g, +)
 
         l = NewLayer4{GRAPH_T}()
@@ -89,7 +89,7 @@
 
         GraphNeuralNetworks.update_node(l::NewLayer5{GRAPH_T}, m̄, xi) = rand(T, 2*out_channel, size(xi, 2))
         GraphNeuralNetworks.update_edge(l::NewLayer5{GRAPH_T}, m, e) = m
-        GraphNeuralNetworks.message(l::NewLayer5{GRAPH_T}, xi, xj, e) = ones(T, out_channel, size(e,2))
+        GraphNeuralNetworks.compute_message(l::NewLayer5{GRAPH_T}, xi, xj, e) = ones(T, out_channel, size(e,2))
         (l::NewLayer5{GRAPH_T})(g) = GraphNeuralNetworks.propagate(l, g, +)
 
         l = NewLayer5{GRAPH_T}()
@@ -110,7 +110,7 @@
 
         NewLayerW(in, out) = NewLayerW{GRAPH_T}(randn(T, out, in))
 
-        GraphNeuralNetworks.message(l::NewLayerW{GRAPH_T}, x_i, x_j, e_ij) = l.weight * x_j
+        GraphNeuralNetworks.compute_message(l::NewLayerW{GRAPH_T}, x_i, x_j, e_ij) = l.weight * x_j
         GraphNeuralNetworks.update_node(l::NewLayerW{GRAPH_T}, m, x) = l.weight * x + m
 
         l = NewLayerW(in_channel, out_channel)