From cf7a1fe4dbe2211ec11d5f73fe43b8024b1462de Mon Sep 17 00:00:00 2001 From: Leni Aniva Date: Sat, 17 Jan 2026 22:03:21 -0800 Subject: [PATCH 1/4] 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/4] 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/4] 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) From 8ec36fec3e548e67c64a5160d9b8d39e8aec52ed Mon Sep 17 00:00:00 2001 From: Leni Aniva Date: Fri, 6 Nov 2026 15:33:03 -0800 Subject: [PATCH 4/4] Add unit test --- CHANGELOG.md | 5 ++++ GNNlib/Project.toml | 2 +- GNNlib/src/layers/conv.jl | 26 +++++++++++++------ GraphNeuralNetworks/test/layers/heteroconv.jl | 22 ++++++++++++++-- 4 files changed, 44 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ffdbc7993..e50f4f39e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and the packages adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). Entries link to the pull request that introduced them. +## GNNlib.jl Unreleased + +**Fixed** +- `GATConv` and `GATv2Conv` no longer fail on graphs (or heterograph relations) without edges for some edge categories, or with zero source/destination nodes: the forward pass crashed reshaping the empty output, and the gradient failed with a `DimensionMismatch`. + ## GNNGraphs.jl 1.6.0 - 2026-11-04 **Added** diff --git a/GNNlib/Project.toml b/GNNlib/Project.toml index a0dd0597e..de495d764 100644 --- a/GNNlib/Project.toml +++ b/GNNlib/Project.toml @@ -1,7 +1,7 @@ name = "GNNlib" uuid = "a6a84749-d869-43f8-aacc-be26a1996e48" authors = ["Carlo Lucibello and contributors"] -version = "1.4.1" +version = "1.4.2" [workspace] projects = ["test", "docs", "../GNNGraphs"] diff --git a/GNNlib/src/layers/conv.jl b/GNNlib/src/layers/conv.jl index fa9952d0d..a8ab59db3 100644 --- a/GNNlib/src/layers/conv.jl +++ b/GNNlib/src/layers/conv.jl @@ -134,10 +134,15 @@ function gat_conv(l, g::AbstractGNNGraph, x, e::Union{Nothing, AbstractMatrix} = # a hand-written message passing message = Fix1(gat_message, l) m = apply_edges(message, g, Wxi, Wxj, e) - α = softmax_edge_neighbors(g, m.logα) - α = dropout(α, l.dropout) - β = α .* m.Wxj - x = aggregate_neighbors(g, +, β) + if isempty(m.logα) + # no edges: Zygote computes a wrong-sized gradient for `α .* m.Wxj` on empty arrays + x = zero(Wxi) + else + α = softmax_edge_neighbors(g, m.logα) + α = dropout(α, l.dropout) + β = α .* m.Wxj + x = aggregate_neighbors(g, +, β) + end if !l.concat x = mean(x, dims = 2) @@ -189,10 +194,15 @@ function gatv2_conv(l, g::AbstractGNNGraph, x, e::Union{Nothing, AbstractMatrix} message = Fix1(gatv2_message, l) m = apply_edges(message, g, Wxi, Wxj, e) - α = softmax_edge_neighbors(g, m.logα) - α = dropout(α, l.dropout) - β = α .* m.Wxj - x = aggregate_neighbors(g, +, β) + if isempty(m.logα) + # no edges: Zygote computes a wrong-sized gradient for `α .* m.Wxj` on empty arrays + x = zero(Wxi) + else + α = softmax_edge_neighbors(g, m.logα) + α = dropout(α, l.dropout) + β = α .* m.Wxj + x = aggregate_neighbors(g, +, β) + end if !l.concat x = mean(x, dims = 2) diff --git a/GraphNeuralNetworks/test/layers/heteroconv.jl b/GraphNeuralNetworks/test/layers/heteroconv.jl index 4a6d40ca3..81c1f19cb 100644 --- a/GraphNeuralNetworks/test/layers/heteroconv.jl +++ b/GraphNeuralNetworks/test/layers/heteroconv.jl @@ -146,10 +146,28 @@ x = (A = rand(Float32, 4, 2), B = rand(Float32, 4, 3)) layers = HeteroGraphConv((:A, :to, :B) => GATv2Conv(4 => 2), (:B, :to, :A) => GATv2Conv(4 => 2)); - y = layers(hg, x); + y = layers(hg, x); @test size(y.A) == (2, 2) && size(y.B) == (2, 3) end - + + @testset "$L with empty relations" for L in (GATConv, GATv2Conv) + g = GNNHeteroGraph(((:A, :to, :B) => ([1, 2], [3, 4]), + (:B, :to, :A) => (Int[], Int[]), + (:C, :to, :A) => (Int[], Int[]), + (:A, :to, :C) => (Int[], Int[])); + num_nodes = Dict(:A => 3, :B => 5, :C => 0)) + x = (A = rand(Float32, 4, 3), B = rand(Float32, 4, 5), C = rand(Float32, 4, 0)) + layers = HeteroGraphConv([et => L(4 => 2; heads = 2, dropout = 0) for et in g.etypes]) + y = layers(g, x) + @test size(y.A) == (4, 3) && size(y.B) == (4, 5) && size(y.C) == (4, 0) + @test y.A == zeros(Float32, 4, 3) + + f(x) = sum(sum, layers(g, x)) + dx = gradient(f, x)[1] + _, (ndx,) = finitediff_withgradient(f, x) + @test dx.A ≈ ndx.A rtol = 1e-3 + end + @testset "GCNConv" begin g = rand_bipartite_heterograph((2,3), 6) x = (A = rand(Float32, 4,2), B = rand(Float32, 4, 3))