From cf7a1fe4dbe2211ec11d5f73fe43b8024b1462de Mon Sep 17 00:00:00 2001 From: Leni Aniva Date: Sat, 17 Jan 2026 22:03:21 -0800 Subject: [PATCH 1/3] fix: Gate empty result in GATConv --- GNNlib/src/layers/conv.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/GNNlib/src/layers/conv.jl b/GNNlib/src/layers/conv.jl index ab40f13f2..5c9f0a856 100644 --- a/GNNlib/src/layers/conv.jl +++ b/GNNlib/src/layers/conv.jl @@ -139,11 +139,12 @@ function gat_conv(l, g::AbstractGNNGraph, x, e::Union{Nothing, AbstractMatrix} = α = dropout(α, l.dropout) β = α .* m.Wxj x = aggregate_neighbors(g, +, β) + width = size(x, 1) if !l.concat x = mean(x, dims = 2) end - x = reshape(x, :, size(x, 3)) # return a matrix + x = reshape(x, width, size(x, 3)) # return a matrix x = l.σ.(x .+ l.bias) return x From 4f9ffeff1b92f56b10e8f6796ccc2e4d317a2dce Mon Sep 17 00:00:00 2001 From: Leni Aniva Date: Sun, 18 Jan 2026 00:44:22 -0800 Subject: [PATCH 2/3] fix: Handle multiple heads --- GNNlib/src/layers/conv.jl | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/GNNlib/src/layers/conv.jl b/GNNlib/src/layers/conv.jl index 5c9f0a856..f5c466834 100644 --- a/GNNlib/src/layers/conv.jl +++ b/GNNlib/src/layers/conv.jl @@ -124,13 +124,10 @@ function gat_conv(l, g::AbstractGNNGraph, x, e::Union{Nothing, AbstractMatrix} = _, chout = l.channel heads = l.heads - Wxi = Wxj = l.dense_x(xj) - Wxi = Wxj = reshape(Wxj, chout, heads, :) - - if xi !== xj - Wxi = l.dense_x(xi) - Wxi = reshape(Wxi, chout, heads, :) - end + Wxj = l.dense_x(xj) + Wxj = reshape(Wxj, chout, heads, :) + Wxi = l.dense_x(xi) + Wxi = reshape(Wxi, chout, heads, :) # a hand-written message passing message = Fix1(gat_message, l) @@ -139,10 +136,12 @@ function gat_conv(l, g::AbstractGNNGraph, x, e::Union{Nothing, AbstractMatrix} = α = dropout(α, l.dropout) β = α .* m.Wxj x = aggregate_neighbors(g, +, β) - width = size(x, 1) if !l.concat x = mean(x, dims = 2) + width = size(x, 1) + else + width = size(x, 1) * size(x, 2) end x = reshape(x, width, size(x, 3)) # return a matrix x = l.σ.(x .+ l.bias) @@ -195,8 +194,11 @@ function gatv2_conv(l, g::AbstractGNNGraph, x, e::Union{Nothing, AbstractMatrix} if !l.concat x = mean(x, dims = 2) + width = size(x, 1) + else + width = size(x, 1) * size(x, 2) end - x = reshape(x, :, size(x, 3)) + x = reshape(x, width, size(x, 3)) x = l.σ.(x .+ l.bias) return x end From a5651b950711a643109b653195fc2a84d1612984 Mon Sep 17 00:00:00 2001 From: Leni Aniva Date: Wed, 2 Sep 2026 23:01:27 -0700 Subject: [PATCH 3/3] fix: Construct stack size --- GNNlib/src/layers/conv.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/GNNlib/src/layers/conv.jl b/GNNlib/src/layers/conv.jl index f5c466834..fa9952d0d 100644 --- a/GNNlib/src/layers/conv.jl +++ b/GNNlib/src/layers/conv.jl @@ -121,13 +121,15 @@ function gat_conv(l, g::AbstractGNNGraph, x, e::Union{Nothing, AbstractMatrix} = g = add_self_loops(g) end + _, n_stackj = size(xj) + _, n_stacki = size(xi) _, chout = l.channel heads = l.heads Wxj = l.dense_x(xj) - Wxj = reshape(Wxj, chout, heads, :) + Wxj = reshape(Wxj, chout, heads, n_stackj) Wxi = l.dense_x(xi) - Wxi = reshape(Wxi, chout, heads, :) + Wxi = reshape(Wxi, chout, heads, n_stacki) # a hand-written message passing message = Fix1(gat_message, l)