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..68203646c 100644 --- a/GNNlib/test/msgpass.jl +++ b/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 diff --git a/GNNlib/test/test_module.jl b/GNNlib/test/test_module.jl index 3478ca9ff..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 @@ -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 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..bf2f58e6a 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 @@ -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 diff --git a/GraphNeuralNetworks/test/layers/temporalconv.jl b/GraphNeuralNetworks/test/layers/temporalconv.jl index 3f04362b8..0f926ff1d 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,15 +155,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(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 @@ -170,9 +173,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 @testitem "DCGRU" setup=[TemporalConvTestModule, TestModule] begin @@ -184,15 +187,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 +204,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..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 @@ -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 @@ -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