Enable Enzyme gradient tests across the test suite (Julia ≥ 1.12)

GNNlib/test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
@@ -22,4 +23,4 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 GNNlib = {path = ".."}
 
 [compat]
-NNlib = "0.9.38"
+NNlib = "0.9.45"

GNNlib/test/msgpass.jl

@@ -128,11 +128,19 @@ end
 
 @testitem "propagate" setup=[TestModuleGNNlib] begin
     using .TestModuleGNNlib
+    using Flux: Flux
 
     @testset "copy_xj +" begin
         for g in TEST_GRAPHS
             f(g, x) = propagate(copy_xj, g, +, xj = x)
-            test_gradients(f, g, g.x; test_grad_f=false)
+            # On :dense graphs the `copy_xj`/`+` specialization multiplies by
+            # `adjacency_matrix(g, T; weighted = false)`, whose result is union-typed
+            # (`to_dense` rebinds `A` across its branches), and Enzyme fails there
+            # with an IllegalTypeAnalysisException. The other storage types are fine.
+            test_gradients(f, g, g.x; test_grad_f=false,
+                           ad_backends = get_graph_type(g) == :dense ?
+                                         [Flux.AutoZygote(), Flux.AutoMooncake()] :
+                                         [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()])
         end
     end
 

GNNlib/test/test_module.jl

@@ -37,7 +37,7 @@ import Reexport: @reexport
 @reexport using SparseArrays
 @reexport using Test, Random, Statistics
 @reexport using MLDataDevices
-using Functors: fmapstructure_with_path
+using Functors: fmap, fmapstructure_with_path
 using FiniteDifferences: FiniteDifferences
 using Zygote: Zygote
 using Flux: Flux
@@ -48,10 +48,18 @@ if TEST_MOONCAKE
     import Mooncake
 end
 
+# Enzyme.jl requires Julia >= 1.12 for the graph code paths exercised here:
+# on 1.10 it recurses forever through `fieldnames(::UnionAll)` in GNNGraph's
+# property sugar (https://github.com/EnzymeAD/Enzyme.jl/issues/3423).
+const TEST_ENZYME = VERSION >= v"1.12"
+if TEST_ENZYME
+    import Enzyme # loads FluxEnzymeExt, which backs Flux.AutoEnzyme()
+end
+
 # from this module
 export D_IN, D_OUT, GRAPH_TYPES, TEST_GRAPHS,
-       test_gradients, finitediff_withgradient, 
-       check_equal_leaves, gpu_backend, TEST_MOONCAKE
+       test_gradients, finitediff_withgradient,
+       check_equal_leaves, gpu_backend, TEST_MOONCAKE, TEST_ENZYME
 
 
 const D_IN = 3
@@ -95,7 +103,8 @@ function test_gradients(
             test_grad_x = true,
             # Reference AD: finite differences for CPU tests, Zygote (on CPU) for GPU tests.
             reference = test_gpu ? Flux.AutoZygote() : :finitediff,
-            ad_backends = test_gpu ? [] : [Flux.AutoZygote(), Flux.AutoMooncake()],
+            ad_backends = test_gpu ? [] :
+                          [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()],
             loss = (f, g, xs...) -> mean(f(g, xs...)),
             )
 
@@ -104,6 +113,11 @@ function test_gradients(
         ad_backends = filter(b -> !(b isa Flux.AutoMooncake), ad_backends)
     end
 
+    # Enzyme requires Julia >= 1.12. Layers it cannot handle opt out per call site.
+    if !TEST_ENZYME
+        ad_backends = filter(b -> !(b isa Flux.AutoEnzyme), ad_backends)
+    end
+
     ## Let's make sure first that the forward pass works.
     l = loss(f, graph, xs...)
     @assert l isa Number
@@ -149,7 +163,10 @@ function test_gradients(
             f64 = f |> Flux.f64
             ps, re = Flux.destructure(f64)
             y_fd, g_fd = finitediff_withgradient(ps -> loss(re(ps), graph, xs...), ps)
-            y_fd, (re(g_fd[1]),)
+            # Rebuild from a zeroed model, so non-trainable arrays (e.g. BatchNorm's
+            # running stats) compare as zero gradients instead of as model values.
+            _, re0 = Flux.destructure(fmap(x -> x isa AbstractArray ? zero(x) : x, f64))
+            y_fd, (re0(g_fd[1]),)
         else
             Flux.withgradient(f -> loss(f, graph, xs...), reference, f)
         end

GraphNeuralNetworks/test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
@@ -9,6 +10,7 @@ GraphNeuralNetworks = "cffab07f-9bc2-4db1-8861-388f63bf7694"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -19,3 +21,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [sources]
 GraphNeuralNetworks = {path = ".."}
+
+[compat]
+NNlib = "0.9.45"

GraphNeuralNetworks/test/layers/basic.jl

@@ -18,7 +18,11 @@
 
         Flux.testmode!(gnn)
 
-        test_gradients(gnn, g, x, rtol = 1e-5, ad_backends = [Flux.AutoZygote()])
+        # Enzyme fails with an IllegalTypeAnalysisException on :dense graphs, where
+        # `propagate(copy_xj, g, +)` multiplies by a union-typed adjacency matrix.
+        test_gradients(gnn, g, x, rtol = 1e-5,
+                       ad_backends = GRAPH_T == :dense ? [Flux.AutoZygote()] :
+                                     [Flux.AutoZygote(), Flux.AutoEnzyme()])
 
         @testset "constructor with names" begin
             m = GNNChain(GCNConv(din => d),
@@ -53,7 +57,10 @@
 
             Flux.trainmode!(gnn)
 
-            test_gradients(gnn, g, x, rtol = 1e-4, atol=1e-4, ad_backends = [Flux.AutoZygote()])
+            # Same :dense limitation as above; Enzyme handles the other graph types.
+            test_gradients(gnn, g, x, rtol = 1e-4, atol=1e-4,
+                           ad_backends = GRAPH_T == :dense ? [Flux.AutoZygote()] :
+                                         [Flux.AutoZygote(), Flux.AutoEnzyme()])
         end
     end
 

GraphNeuralNetworks/test/layers/conv.jl

@@ -86,7 +86,8 @@ end
     for g in TEST_GRAPHS
         g = add_self_loops(g)
         @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
-        test_gradients(l, g, g.x, rtol = RTOL_LOW, ad_backends = [Flux.AutoZygote()])
+        test_gradients(l, g, g.x, rtol = RTOL_LOW,
+                       ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()])
     end
 
     @testset "bias=false" begin
@@ -124,7 +125,12 @@ end
     l = GraphConv(D_IN => D_OUT)
     for g in TEST_GRAPHS
         @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
-        test_gradients(l, g, g.x, rtol = RTOL_HIGH)
+        # Enzyme fails with an IllegalTypeAnalysisException on :dense graphs, where
+        # `propagate(copy_xj, g, +)` multiplies by a union-typed adjacency matrix.
+        test_gradients(l, g, g.x, rtol = RTOL_HIGH,
+                       ad_backends = get_graph_type(g) == :dense ?
+                                     [Flux.AutoZygote(), Flux.AutoMooncake()] :
+                                     [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()])
     end
 
     l = GraphConv(D_IN => D_OUT, tanh, bias = false, aggr = mean)
@@ -236,10 +242,14 @@ end
     num_layers = 3
     l = GatedGraphConv(D_OUT, num_layers)
     @test size(l.weight) == (D_OUT, D_OUT, num_layers)
-
     for g in TEST_GRAPHS
         @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
-        test_gradients(l, g, g.x, rtol = RTOL_HIGH)
+        # Enzyme fails with an IllegalTypeAnalysisException on :dense graphs, where
+        # `propagate(copy_xj, g, +)` multiplies by a union-typed adjacency matrix.
+        test_gradients(l, g, g.x, rtol = RTOL_HIGH,
+                       ad_backends = get_graph_type(g) == :dense ?
+                                     [Flux.AutoZygote(), Flux.AutoMooncake()] :
+                                     [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()])
     end
 end
 
@@ -332,7 +342,12 @@ end
     l = SAGEConv(D_IN => D_OUT, tanh, bias = false, aggr = +)
     for g in TEST_GRAPHS
         @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
-        test_gradients(l, g, g.x, rtol = RTOL_HIGH)
+        # Enzyme fails with an IllegalTypeAnalysisException on :dense graphs, where
+        # `propagate(copy_xj, g, +)` multiplies by a union-typed adjacency matrix.
+        test_gradients(l, g, g.x, rtol = RTOL_HIGH,
+                       ad_backends = get_graph_type(g) == :dense ?
+                                     [Flux.AutoZygote(), Flux.AutoMooncake()] :
+                                     [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()])
     end
 end
 
@@ -433,8 +448,8 @@ end
         @test size(y[1]) == (D_OUT, g.num_nodes)
         @test size(y[2]) == (D_OUT, g.num_edges)
         function loss(l, g, x, e)
-            y = l(g, x, e)
-            return mean(y[1]) + sum(y[2])
+            ŷ = l(g, x, e)
+            return mean(ŷ[1]) + sum(ŷ[2])
         end
         test_gradients(l, g, g.x, g.e, rtol = RTOL_LOW; loss)
     end
@@ -484,18 +499,24 @@ end
 
 @testitem "SGConv" setup=[TolSnippet, TestModule] begin
     using .TestModule
-    K = [1, 2, 3] # for different number of hops       
+    K = [1, 2, 3] # for different number of hops
+    # Enzyme is skipped: it fails with an IllegalTypeAnalysisException on every
+    # storage type, because `degree(g, T; edge_weight = l.use_edge_weight)` returns
+    # a value whose type is only known at runtime (`Union{Nothing, AbstractVector}`)
+    # and the k-hop loop propagates that union through each `propagate` call.
     for k in K
         l = SGConv(D_IN => D_OUT, k, add_self_loops = true)
         for g in TEST_GRAPHS
             @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
-            test_gradients(l, g, g.x, rtol = RTOL_HIGH)
+            test_gradients(l, g, g.x, rtol = RTOL_HIGH,
+                           ad_backends = [Flux.AutoZygote(), Flux.AutoMooncake()])
         end
 
         l = SGConv(D_IN => D_OUT, k, add_self_loops = true)
         for g in TEST_GRAPHS
             @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
-            test_gradients(l, g, g.x, rtol = RTOL_HIGH)
+            test_gradients(l, g, g.x, rtol = RTOL_HIGH,
+                           ad_backends = [Flux.AutoZygote(), Flux.AutoMooncake()])
         end
     end
 end
@@ -514,17 +535,26 @@ end
 @testitem "TAGConv" setup=[TolSnippet, TestModule] begin
     using .TestModule
     K = [1, 2, 3]
+    # Enzyme runs only on :sparse graphs: the other storage types hit the same
+    # union-typed `degree(g, T; edge_weight = l.use_edge_weight)` problem as SGConv,
+    # while Enzyme's type analysis happens to recover on :sparse.
     for k in K
         l = TAGConv(D_IN => D_OUT, k, add_self_loops = true)
         for g in TEST_GRAPHS
             @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
-            test_gradients(l, g, g.x, rtol = RTOL_HIGH)
+            test_gradients(l, g, g.x, rtol = RTOL_HIGH,
+                           ad_backends = get_graph_type(g) == :sparse ?
+                                         [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] :
+                                         [Flux.AutoZygote(), Flux.AutoMooncake()])
         end
 
         l = TAGConv(D_IN => D_OUT, k, add_self_loops = true)
         for g in TEST_GRAPHS
             @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
-            test_gradients(l, g, g.x, rtol = RTOL_HIGH)
+            test_gradients(l, g, g.x, rtol = RTOL_HIGH,
+                           ad_backends = get_graph_type(g) == :sparse ?
+                                         [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] :
+                                         [Flux.AutoZygote(), Flux.AutoMooncake()])
         end
     end
 end
@@ -609,7 +639,7 @@ end
 
 @testitem "DConv" setup=[TolSnippet, TestModule] begin
     using .TestModule
-    K = [1, 2, 3] # for different number of hops
+    K = [1, 2, 3] # for different number of hops       
     for k in K
         l = DConv(D_IN => D_OUT, k)
         for g in TEST_GRAPHS