diff --git a/CHANGELOG.md b/CHANGELOG.md index 2ae90818b..2c8edfd7d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,6 +37,7 @@ Entries link to the pull request that introduced them. **Fixed** - Fixed `Enzyme.gradient` failing when differentiating `GCNConv`, `SGConv` and `TAGConv` on `:dense`/`:sparse` adjacency graphs: their adjacency-matrix fallbacks now convert via the Enzyme-differentiable `GNNGraphs._to_coo_graph` instead of the keyword `GNNGraph` constructor. Requires GNNGraphs ≥ 1.5.2 and, for Enzyme, Enzyme ≥ 0.13.197 ([#703]). +- `GATConv` and `GATv2Conv` skip attention dropout when `dropout == 0`, which also makes them differentiable with Mooncake on CUDA ([#XXX]). ## GraphNeuralNetworks.jl — Unreleased (towards 1.1.1) @@ -238,4 +239,5 @@ Lux implementations of the graph convolutional, pooling, and temporal layers [#704]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/704 [#707]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/707 [#706]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/706 +[#709]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/709 [FluxML/Zygote.jl#1662]: https://github.com/FluxML/Zygote.jl/issues/1662 diff --git a/GNNlib/src/layers/conv.jl b/GNNlib/src/layers/conv.jl index ab40f13f2..fbdf6cc5d 100644 --- a/GNNlib/src/layers/conv.jl +++ b/GNNlib/src/layers/conv.jl @@ -136,7 +136,8 @@ function gat_conv(l, g::AbstractGNNGraph, x, e::Union{Nothing, AbstractMatrix} = message = Fix1(gat_message, l) m = apply_edges(message, g, Wxi, Wxj, e) α = softmax_edge_neighbors(g, m.logα) - α = dropout(α, l.dropout) + # Skip at p == 0: NNlib.dropout fetches the (CUDA) RNG before checking p, which Mooncake cannot trace. + iszero(l.dropout) || (α = dropout(α, l.dropout)) β = α .* m.Wxj x = aggregate_neighbors(g, +, β) @@ -188,7 +189,7 @@ 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) + iszero(l.dropout) || (α = dropout(α, l.dropout)) β = α .* m.Wxj x = aggregate_neighbors(g, +, β)