fix: Gate empty result in GATConv

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**

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"]

GNNlib/src/layers/conv.jl

@@ -121,29 +121,36 @@ 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
 
-    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, n_stackj)
+    Wxi = l.dense_x(xi)
+    Wxi = reshape(Wxi, chout, heads, n_stacki)
 
     # 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)
+        width = size(x, 1)
+    else
+        width = size(x, 1) * size(x, 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
@@ -187,15 +194,23 @@ 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)
+        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

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))