From 0a81391a3a2c46c5997feda63ae90edc7a0e75ec Mon Sep 17 00:00:00 2001 From: Parvm1102 Date: Fri, 28 Aug 2026 13:59:06 +0530 Subject: [PATCH 1/3] =?UTF-8?q?Enable=20Enzyme=20gradient=20tests=20across?= =?UTF-8?q?=20the=20test=20suite=20(Julia=20=E2=89=A5=201.12)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Parvm1102 --- GNNlib/test/Project.toml | 3 +- GNNlib/test/msgpass.jl | 11 +++- GNNlib/test/test_module.jl | 20 ++++++- GraphNeuralNetworks/test/Project.toml | 5 ++ GraphNeuralNetworks/test/layers/basic.jl | 11 +++- GraphNeuralNetworks/test/layers/conv.jl | 56 +++++++++++++++---- .../test/layers/temporalconv.jl | 50 ++++++++++------- GraphNeuralNetworks/test/test_module.jl | 20 ++++++- 8 files changed, 135 insertions(+), 41 deletions(-) diff --git a/GNNlib/test/Project.toml b/GNNlib/test/Project.toml index 6df94bff1..d0f15446d 100644 --- a/GNNlib/test/Project.toml +++ b/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" diff --git a/GNNlib/test/msgpass.jl b/GNNlib/test/msgpass.jl index 12c207507..b2e92a1f8 100644 --- a/GNNlib/test/msgpass.jl +++ b/GNNlib/test/msgpass.jl @@ -128,11 +128,20 @@ end @testitem "propagate" setup=[TestModuleGNNlib] begin using .TestModuleGNNlib + using Flux: Flux + + # 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. + enzyme_backends(g) = get_graph_type(g) == :dense ? + [Flux.AutoZygote(), Flux.AutoMooncake()] : + [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] @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) + test_gradients(f, g, g.x; test_grad_f=false, ad_backends = enzyme_backends(g)) end end diff --git a/GNNlib/test/test_module.jl b/GNNlib/test/test_module.jl index 3478ca9ff..176092b9a 100644 --- a/GNNlib/test/test_module.jl +++ b/GNNlib/test/test_module.jl @@ -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 diff --git a/GraphNeuralNetworks/test/Project.toml b/GraphNeuralNetworks/test/Project.toml index a37e4e330..1a7ab540f 100644 --- a/GraphNeuralNetworks/test/Project.toml +++ b/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" diff --git a/GraphNeuralNetworks/test/layers/basic.jl b/GraphNeuralNetworks/test/layers/basic.jl index 27a39c3a8..b6e93b1d5 100644 --- a/GraphNeuralNetworks/test/layers/basic.jl +++ b/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 diff --git a/GraphNeuralNetworks/test/layers/conv.jl b/GraphNeuralNetworks/test/layers/conv.jl index 7585b817d..5f4843bf2 100644 --- a/GraphNeuralNetworks/test/layers/conv.jl +++ b/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 @@ -121,10 +122,17 @@ end @testitem "GraphConv" setup=[TolSnippet, TestModule] begin using .TestModule + + # Enzyme fails with an IllegalTypeAnalysisException on :dense graphs, where + # `propagate(copy_xj, g, +)` multiplies by a union-typed adjacency matrix. + enzyme_backends(g) = get_graph_type(g) == :dense ? + [Flux.AutoZygote(), Flux.AutoMooncake()] : + [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] + 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) + test_gradients(l, g, g.x, rtol = RTOL_HIGH, ad_backends = enzyme_backends(g)) end l = GraphConv(D_IN => D_OUT, tanh, bias = false, aggr = mean) @@ -237,9 +245,15 @@ end l = GatedGraphConv(D_OUT, num_layers) @test size(l.weight) == (D_OUT, D_OUT, num_layers) + # Enzyme fails with an IllegalTypeAnalysisException on :dense graphs, where + # `propagate(copy_xj, g, +)` multiplies by a union-typed adjacency matrix. + enzyme_backends(g) = get_graph_type(g) == :dense ? + [Flux.AutoZygote(), Flux.AutoMooncake()] : + [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] + 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 = enzyme_backends(g)) end end @@ -329,10 +343,16 @@ end l = SAGEConv(D_IN => D_OUT) @test l.aggr == mean + # Enzyme fails with an IllegalTypeAnalysisException on :dense graphs, where + # `propagate(copy_xj, g, +)` multiplies by a union-typed adjacency matrix. + enzyme_backends(g) = get_graph_type(g) == :dense ? + [Flux.AutoZygote(), Flux.AutoMooncake()] : + [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] + 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) + test_gradients(l, g, g.x, rtol = RTOL_HIGH, ad_backends = enzyme_backends(g)) end end @@ -433,8 +453,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 @@ -485,17 +505,24 @@ end @testitem "SGConv" setup=[TolSnippet, TestModule] begin using .TestModule K = [1, 2, 3] # for different number of hops + + # Enzyme fails with an IllegalTypeAnalysisException on every storage type: + # `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. + ad_backends = [Flux.AutoZygote(), Flux.AutoMooncake()] + 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) 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) end end end @@ -514,17 +541,24 @@ end @testitem "TAGConv" setup=[TolSnippet, TestModule] begin using .TestModule K = [1, 2, 3] + + # Same union-typed `degree(g, T; edge_weight = l.use_edge_weight)` problem as + # SGConv, except that Enzyme's type analysis happens to recover on :sparse. + enzyme_backends(g) = get_graph_type(g) == :sparse ? + [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] : + [Flux.AutoZygote(), Flux.AutoMooncake()] + 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 = enzyme_backends(g)) 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 = enzyme_backends(g)) end end end @@ -609,7 +643,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 diff --git a/GraphNeuralNetworks/test/layers/temporalconv.jl b/GraphNeuralNetworks/test/layers/temporalconv.jl index 3f04362b8..4dd18d2e4 100644 --- a/GraphNeuralNetworks/test/layers/temporalconv.jl +++ b/GraphNeuralNetworks/test/layers/temporalconv.jl @@ -30,7 +30,9 @@ end # sigmoid gates (`Dense(_, sigmoid)`) on a fraction of inputs, so the check is # flaky (Zygote and finite differences agree). Upstream Mooncake bug with # NNlib.sigmoid: https://github.com/chalk-lab/Mooncake.jl/issues/1257. - # Reference against Zygote only, matching the other temporal cells. + # Enzyme is also skipped: compiling the gradient of the packaged cell hangs + # (25+ minutes with no output), so enabling it would stall the whole suite. + # Reference against Zygote only. ad_backends = [Flux.AutoZygote()] # Test with default activation function @@ -62,7 +64,8 @@ end # Mooncake is skipped here (see the TGCNCell test item and # https://github.com/chalk-lab/Mooncake.jl/issues/1257): it returns wrong - # gradients through the cell's sigmoid gates on some inputs. Zygote only. + # gradients through the cell's sigmoid gates on some inputs. Enzyme is also + # skipped: it inherits TGCNCell's compile hang. Zygote only. ad_backends = [Flux.AutoZygote()] # Test with default activation function @@ -105,9 +108,9 @@ end @test size(h) == (out_channel, g.num_nodes) @test size(c) == (out_channel, g.num_nodes) # with no initial state - test_gradients(cell, g, g.x, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(cell, g, g.x, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) # with initial state - test_gradients(cell, g, g.x, (h, c), loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(cell, g, g.x, (h, c), loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) end @testitem "GConvLSTM" setup=[TemporalConvTestModule, TestModule] begin @@ -119,15 +122,15 @@ end y = layer(g, x) @test size(y) == (out_channel, timesteps, g.num_nodes) # with no initial state - test_gradients(layer, g, x, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(layer, g, x, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) # with initial state - test_gradients(layer, g, x, state0, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(layer, g, x, state0, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) # interplay with GNNChain model = GNNChain(GConvLSTM(in_channel => out_channel, 2), Dense(out_channel, 1)) y = model(g, x) @test size(y) == (1, timesteps, g.num_nodes) - test_gradients(model, g, x, rtol = RTOL_LOW, atol = ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(model, g, x, rtol = RTOL_LOW, atol = ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) end @testitem "GConvGRUCell" setup=[TemporalConvTestModule, TestModule] begin @@ -137,9 +140,9 @@ end @test y === h @test size(h) == (out_channel, g.num_nodes) # with no initial state - test_gradients(cell, g, g.x, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(cell, g, g.x, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) # with initial state - test_gradients(cell, g, g.x, h, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(cell, g, g.x, h, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) end @@ -152,27 +155,29 @@ end y = layer(g, x) @test size(y) == (out_channel, timesteps, g.num_nodes) # with no initial state - test_gradients(layer, g, x, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(layer, g, x, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) # with initial state - test_gradients(layer, g, x, state0, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(layer, g, x, state0, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) # interplay with GNNChain model = GNNChain(GConvGRU(in_channel => out_channel, 2), Dense(out_channel, 1)) y = model(g, x) @test size(y) == (1, timesteps, g.num_nodes) - test_gradients(model, g, x, rtol = RTOL_LOW, atol = ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(model, g, x, rtol = RTOL_LOW, atol = ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) end @testitem "DCGRUCell" setup=[TemporalConvTestModule, TestModule] begin using .TemporalConvTestModule, .TestModule + ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()] + cell = DCGRUCell(in_channel => out_channel, 2) y, h = cell(g, g.x) @test y === h @test size(h) == (out_channel, g.num_nodes) # with no initial state - test_gradients(cell, g, g.x, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(cell, g, g.x, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW; ad_backends) # with initial state - test_gradients(cell, g, g.x, h, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(cell, g, g.x, h, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW; ad_backends) end @testitem "DCGRU" setup=[TemporalConvTestModule, TestModule] begin @@ -184,15 +189,15 @@ end y = layer(g, x) @test size(y) == (out_channel, timesteps, g.num_nodes) # with no initial state - test_gradients(layer, g, x, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(layer, g, x, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) # with initial state - test_gradients(layer, g, x, state0, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(layer, g, x, state0, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) # interplay with GNNChain model = GNNChain(DCGRU(in_channel => out_channel, 2), Dense(out_channel, 1)) y = model(g, x) @test size(y) == (1, timesteps, g.num_nodes) - test_gradients(model, g, x, rtol = RTOL_LOW, atol = ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(model, g, x, rtol = RTOL_LOW, atol = ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) end @testitem "EvolveGCNOCell" setup=[TemporalConvTestModule, TestModule] begin @@ -201,9 +206,14 @@ end y, state = cell(g, g.x) @test size(y) == (out_channel, g.num_nodes) # with no initial state - test_gradients(cell, g, g.x, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) - # with initial state - test_gradients(cell, g, g.x, state, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote()]) + test_gradients(cell, g, g.x, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, + ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) + # with initial state: the cell returns the same array in the output and in the + # state, and Enzyme accumulates both paths into the one shadow that aliased + # primal memory must share (EnzymeAD/Enzyme.jl#3408), so it disagrees with the + # per-slot partials Zygote and finite differences report. + test_gradients(cell, g, g.x, state, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, + ad_backends = [Flux.AutoZygote()]) end @testitem "EvolveGCNO" setup=[TemporalConvTestModule, TestModule] begin diff --git a/GraphNeuralNetworks/test/test_module.jl b/GraphNeuralNetworks/test/test_module.jl index de5e714d1..f4a60780c 100644 --- a/GraphNeuralNetworks/test/test_module.jl +++ b/GraphNeuralNetworks/test/test_module.jl @@ -35,6 +35,14 @@ 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 Base export mean, randn, SparseArrays, AbstractSparseMatrix @@ -49,8 +57,8 @@ export random_regular_graph, erdos_renyi # 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 From 40701f6352a1ae361cb3ddcf58487c74859e96d9 Mon Sep 17 00:00:00 2001 From: Parvm1102 Date: Fri, 28 Aug 2026 14:12:58 +0530 Subject: [PATCH 2/3] Passing the ad_backends directly Signed-off-by: Parvm1102 --- GNNlib/test/msgpass.jl | 17 ++-- GraphNeuralNetworks/test/layers/conv.jl | 80 +++++++++---------- .../test/layers/temporalconv.jl | 6 +- 3 files changed, 48 insertions(+), 55 deletions(-) diff --git a/GNNlib/test/msgpass.jl b/GNNlib/test/msgpass.jl index b2e92a1f8..68203646c 100644 --- a/GNNlib/test/msgpass.jl +++ b/GNNlib/test/msgpass.jl @@ -130,18 +130,17 @@ end using .TestModuleGNNlib using Flux: Flux - # 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. - enzyme_backends(g) = get_graph_type(g) == :dense ? - [Flux.AutoZygote(), Flux.AutoMooncake()] : - [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] - @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, ad_backends = enzyme_backends(g)) + # 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 diff --git a/GraphNeuralNetworks/test/layers/conv.jl b/GraphNeuralNetworks/test/layers/conv.jl index 5f4843bf2..bf2f58e6a 100644 --- a/GraphNeuralNetworks/test/layers/conv.jl +++ b/GraphNeuralNetworks/test/layers/conv.jl @@ -122,17 +122,15 @@ end @testitem "GraphConv" setup=[TolSnippet, TestModule] begin using .TestModule - - # Enzyme fails with an IllegalTypeAnalysisException on :dense graphs, where - # `propagate(copy_xj, g, +)` multiplies by a union-typed adjacency matrix. - enzyme_backends(g) = get_graph_type(g) == :dense ? - [Flux.AutoZygote(), Flux.AutoMooncake()] : - [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] - 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, ad_backends = enzyme_backends(g)) + # 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) @@ -244,16 +242,14 @@ end num_layers = 3 l = GatedGraphConv(D_OUT, num_layers) @test size(l.weight) == (D_OUT, D_OUT, num_layers) - - # Enzyme fails with an IllegalTypeAnalysisException on :dense graphs, where - # `propagate(copy_xj, g, +)` multiplies by a union-typed adjacency matrix. - enzyme_backends(g) = get_graph_type(g) == :dense ? - [Flux.AutoZygote(), Flux.AutoMooncake()] : - [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] - 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, ad_backends = enzyme_backends(g)) + # 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 @@ -343,16 +339,15 @@ end l = SAGEConv(D_IN => D_OUT) @test l.aggr == mean - # Enzyme fails with an IllegalTypeAnalysisException on :dense graphs, where - # `propagate(copy_xj, g, +)` multiplies by a union-typed adjacency matrix. - enzyme_backends(g) = get_graph_type(g) == :dense ? - [Flux.AutoZygote(), Flux.AutoMooncake()] : - [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] - 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, ad_backends = enzyme_backends(g)) + # 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 @@ -504,25 +499,24 @@ end @testitem "SGConv" setup=[TolSnippet, TestModule] begin using .TestModule - K = [1, 2, 3] # for different number of hops - - # Enzyme fails with an IllegalTypeAnalysisException on every storage type: - # `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. - ad_backends = [Flux.AutoZygote(), Flux.AutoMooncake()] - + 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; ad_backends) + 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; ad_backends) + test_gradients(l, g, g.x, rtol = RTOL_HIGH, + ad_backends = [Flux.AutoZygote(), Flux.AutoMooncake()]) end end end @@ -541,24 +535,26 @@ end @testitem "TAGConv" setup=[TolSnippet, TestModule] begin using .TestModule K = [1, 2, 3] - - # Same union-typed `degree(g, T; edge_weight = l.use_edge_weight)` problem as - # SGConv, except that Enzyme's type analysis happens to recover on :sparse. - enzyme_backends(g) = get_graph_type(g) == :sparse ? - [Flux.AutoZygote(), Flux.AutoMooncake(), Flux.AutoEnzyme()] : - [Flux.AutoZygote(), Flux.AutoMooncake()] - + # 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, ad_backends = enzyme_backends(g)) + 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, ad_backends = enzyme_backends(g)) + 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 diff --git a/GraphNeuralNetworks/test/layers/temporalconv.jl b/GraphNeuralNetworks/test/layers/temporalconv.jl index 4dd18d2e4..0f926ff1d 100644 --- a/GraphNeuralNetworks/test/layers/temporalconv.jl +++ b/GraphNeuralNetworks/test/layers/temporalconv.jl @@ -168,16 +168,14 @@ end @testitem "DCGRUCell" setup=[TemporalConvTestModule, TestModule] begin using .TemporalConvTestModule, .TestModule - ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()] - cell = DCGRUCell(in_channel => out_channel, 2) y, h = cell(g, g.x) @test y === h @test size(h) == (out_channel, g.num_nodes) # with no initial state - test_gradients(cell, g, g.x, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW; ad_backends) + test_gradients(cell, g, g.x, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) # with initial state - test_gradients(cell, g, g.x, h, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW; ad_backends) + test_gradients(cell, g, g.x, h, loss=cell_loss, rtol=RTOL_LOW, atol=ATOL_LOW, ad_backends = [Flux.AutoZygote(), Flux.AutoEnzyme()]) end @testitem "DCGRU" setup=[TemporalConvTestModule, TestModule] begin From 3fde01e04d6fc4aaf29efac055618b40609d1d9f Mon Sep 17 00:00:00 2001 From: Parvm1102 Date: Fri, 28 Aug 2026 14:59:30 +0530 Subject: [PATCH 3/3] Zero non-trainable leaves in the finite-difference reference `Flux.destructure`'s `re` writes the gradient into the trainable leaves but keeps every other array (e.g. BatchNorm's running stats) at its model value. Zygote returns `nothing` for those fields so `check_equal_leaves` skipped them, but Enzyme returns a full shadow model with the correct zero gradient there and the comparison failed on state that is not a gradient at all. Rebuild the reference from a zeroed model so those leaves compare as zeros. This enables Enzyme on the GNNChain Parallel/BatchNorm test site for non-dense graphs; verified neutral for EdgeConv and GNNlib's propagate. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Parvm1102 --- GNNlib/test/test_module.jl | 7 +++++-- GraphNeuralNetworks/test/test_module.jl | 7 +++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/GNNlib/test/test_module.jl b/GNNlib/test/test_module.jl index 176092b9a..1af8c5a2f 100644 --- a/GNNlib/test/test_module.jl +++ b/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 @@ -163,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 diff --git a/GraphNeuralNetworks/test/test_module.jl b/GraphNeuralNetworks/test/test_module.jl index f4a60780c..d48309ef5 100644 --- a/GraphNeuralNetworks/test/test_module.jl +++ b/GraphNeuralNetworks/test/test_module.jl @@ -23,7 +23,7 @@ using GraphNeuralNetworks using Test using Statistics, Random using Flux -using Functors: fmapstructure_with_path +using Functors: fmap, fmapstructure_with_path using Graphs using ChainRulesTestUtils, FiniteDifferences using Zygote: Zygote @@ -163,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