# Best options for Graph Neural Network integration with Flux

**URL:** <https://discourse.julialang.org/t/best-options-for-graph-neural-network-integration-with-flux/133368>\
**Category:** Machine Learning\
**Tags:** package, flux, machine-learning, geometricflux, graphneuralnetworks\
**Created:** [October 22, 2025, 10:59pm UTC](https://discourse.julialang.org/t/best-options-for-graph-neural-network-integration-with-flux/133368 "2025-10-22T22:59:04Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![AriMarkowitz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/arimarkowitz/32/216087_2.png) [@AriMarkowitz](https://discourse.julialang.org/u/AriMarkowitz)\
**Post date:** [October 22, 2025, 10:59pm UTC](https://discourse.julialang.org/t/best-options-for-graph-neural-network-integration-with-flux/133368/1 "2025-10-22T22:59:04Z")

</div>

Hello!

I built a library that utilizes GraphNeuralNetworks.jl and Flux.jl, but the GraphNeuralNetworks.jl tests are currently failing in Julia 1.12 (and the same failure is causing my custom library to crash when used). I am wondering if anyone has had similar problems, or if anyone has suggestions on alternatives that are compatible with Flux. I am also aware of GeometricFlux.jl but I was under the impression that GraphNeuralNetworks.jl is the best maintained library in this domain that is compatible with Flux.

Thanks for your help!

---

<div class="post-metadata">

**Author:** ![AriMarkowitz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/arimarkowitz/32/216087_2.png) [@AriMarkowitz](https://discourse.julialang.org/u/AriMarkowitz)\
**Post date:** [October 22, 2025, 11:20pm UTC](https://discourse.julialang.org/t/best-options-for-graph-neural-network-integration-with-flux/133368/2 "2025-10-22T23:20:42Z")

</div>

[FYI I submitted a Github issue regarding the test failures here.](https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/issues/623)

---

<div class="post-metadata">

**Author:** ![monty](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/monty/32/44234_2.png) [@monty](https://discourse.julialang.org/u/monty)\
**Post date:** [October 23, 2025, 3:00pm UTC](https://discourse.julialang.org/t/best-options-for-graph-neural-network-integration-with-flux/133368/3 "2025-10-23T15:00:25Z")

</div>

I am also very curious how we can resolve this issues. Any suggestions will be highly appreciated. Thank you!

---

<div class="post-metadata">

**Author:** ![CarloLucibello](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/carlolucibello/32/3278_2.png) [@CarloLucibello](https://discourse.julialang.org/u/CarloLucibello)\
**Post date:** [October 23, 2025, 8:47pm UTC](https://discourse.julialang.org/t/best-options-for-graph-neural-network-integration-with-flux/133368/4 "2025-10-23T20:47:12Z")

</div>

The problem is upstream, Zygote does not support v1.12 yet.

> <https://github.com/FluxML/Zygote.jl/issues/1580>
>
> For example:https://s3.amazonaws.com/julialang-reports/nanosoldier/pkgeval/by\_ha…sh/c37ee51\_vs\_bd445fa/JOLI.primary.log
> 
> \`\`\`
> ErrorException("IR verification failed.\\n Code location: /home/pkgeval/.julia/packages/JOLI/2dSnd/test/test\_rrules.jl:8\\n Method instance: MethodInstance for ZygoteRules.\_pullback(::Zygote.Context{false}, ::var\\"#497#498\\", ::Vector{Float32})")
> error at ./error.jl:54
> unknown function (ip: 0x79a739f12f9a) at (unknown file)
> \_jl\_invoke at /source/src/gf.c:3696 \[inlined\]
> ijl\_apply\_generic at /source/src/gf.c:3896
> jl\_apply at /source/src/julia.h:2372 \[inlined\]
> jl\_f\_invokelatest at /source/src/builtins.c:875
> \_jl\_invoke at /source/src/gf.c:3696 \[inlined\]
> ijl\_apply\_generic at /source/src/gf.c:3896
> jl\_apply at /source/src/julia.h:2372 \[inlined\]
> jl\_f\_\_apply\_iterate at /source/src/builtins.c:862
> raise\_error at ./../usr/share/julia/Compiler/src/ssair/verify.jl:125
> jfptr\_raise\_error\_54024.1 at /opt/julia/lib/julia/sys.so (unknown line)
> \_jl\_invoke at /source/src/gf.c:3696 \[inlined\]
> ijl\_apply\_generic at /source/src/gf.c:3896
> check\_op at ./../usr/share/julia/Compiler/src/ssair/verify.jl:71
> verify\_ir at ./../usr/share/julia/Compiler/src/ssair/verify.jl:429
> run\_passes\_ipo\_safe at ./../usr/share/julia/Compiler/src/optimize.jl:1043
> run\_passes\_ipo\_safe at ./../usr/share/julia/Compiler/src/optimize.jl:1027 \[inlined\]
> optimize at ./../usr/share/julia/Compiler/src/optimize.jl:1002
> jfptr\_optimize\_121107.1 at /opt/julia/lib/julia/sys.so (unknown line)
> \_jl\_invoke at /source/src/gf.c:3696 \[inlined\]
> \`\`\`

I don’t know if anyone is working on the issue.

---

<div class="post-metadata">

**Author:** ![AriMarkowitz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/arimarkowitz/32/216087_2.png) [@AriMarkowitz](https://discourse.julialang.org/u/AriMarkowitz)\
**Post date:** [October 23, 2025, 9:04pm UTC](https://discourse.julialang.org/t/best-options-for-graph-neural-network-integration-with-flux/133368/5 "2025-10-23T21:04:14Z")

</div>

Thank you for the response @CarloLucibello. Do you think I may be able to achieve a workaround by making our GraphNeuralNetwork dependent model explicitly use Enzyme.jl?

---

<div class="post-metadata">

**Author:** ![CarloLucibello](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/carlolucibello/32/3278_2.png) [@CarloLucibello](https://discourse.julialang.org/u/CarloLucibello)\
**Post date:** [October 24, 2025, 7:15am UTC](https://discourse.julialang.org/t/best-options-for-graph-neural-network-integration-with-flux/133368/6 "2025-10-24T07:15:42Z")

</div>

My suggestion is to stick with julia v1.11 for the time being. Enzyme is not officially compatible and likely it is not going to work:

> <https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/issues/389>
>
> Here we will keep track of compatibility with Enzyme for taking gradients. 
> Fir…st think is to collect a few examples to run.

> <https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/issues/609>
>
> Trying to reproduce \[bug 2472 on Enzyme\](https://github.com/EnzymeAD/Enzyme.jl/i…ssues/2472), I kept reducing my examples until I found that the most trivial use of \`GCNConv\` results in Enzyme throwing an error about "Constant memory is stored (or returned) to a differentiable variable", and sometimes even in segfaults. Here's the example:
> \`\`\`julia
> using Lux, GNNLux, Random, Enzyme, Optimisers
> 
> Lux.@concrete struct GraphPolicyLux \<: GNNContainerLayer{(:conv1, :conv2, :dense)}
> conv1
> conv2
> dense
> end
> 
> function GraphPolicyLux(nin::Int, d::Int, n\_nodes::Int)
> conv1 = GCNConv(nin =\> d)
> conv2 = GCNConv(d =\> d)
> dense = Dense(d =\> n\_nodes, sigmoid)
> return GraphPolicyLux(conv1, conv2, dense)
> end
> 
> function (model::GraphPolicyLux)(g::GNNGraph, x, ps, st)
> dense = StatefulLuxLayer{true}(model.dense, ps.dense, GNNLux.\_getstate(st, :dense))
> x, st\_c1 = model.conv1(g, x, ps.conv1, st.conv1)
> x = tanh.(x)
> x, st\_c2 = model.conv2(g, x, ps.conv2, st.conv2)
> x = relu.(x)
> new\_weights = dense(x)
> return new\_weights, (conv1 = st\_c1, conv2 = st\_c2)
> end
> 
> 
> \# Dummy graph construction
> n = 100
> D = 2
> A = sprand(Float32, n, n, 0.2)
> X = rand(Float32, D, n)
> g = GNNGraph(A; ndata=(features = X))
> \# Setup model and parameters
> model = GraphPolicyLux(2, 50, g.num\_nodes)
> ps, st = Lux.setup(Random.default\_rng(), model)
> train\_state = Training.TrainState(model, ps, st, Adam(0.01f0))
> 
> function compute\_loss(model, ps, st, (g, X))
> prop\_A, st = model(g, X, ps, st)
> Reward = 0.0f0
> # This is where a simulation with A prop\_A would go
> for t in 1:100
> Reward -= t \* 0.01
> # Reward -= rand(Float32)
> end
> return Reward, st, 0
> end
> 
> function train\_trigger\_bug!(model, ps, st, g, nepochs::Int=100)
> tstate = Training.TrainState(model, ps, st, Adam(0.01f0))
> data = (g, g.ndata.x)
> for epoch in 1:nepochs
> \_, loss, \_, tstate = Training.single\_train\_step!(AutoEnzyme(), compute\_loss, (g, g.ndata.x), tstate)
> @show epoch
> end
> end
> \`\`\`
> 
> To make sure I wasn't doing something wrong, I copy-pasted the (only) \[example\](https://juliagraphs.org/GraphNeuralNetworks.jl/docs/GNNLux.jl/stable/tutorials/gnn\_intro/) with Lux.jl. This problem is not limited to Lux.jl, though, and the same issue occurs when using Flux.jl. By changing a single line from
> \`\`\`julia
> \_, loss, \_, train\_state = Lux.Training.single\_train\_step!(AutoZygote(), custom\_loss,(g, g.x, g.y), train\_state)
> \`\`\`
> to
> \`\`\`julia
> \_, loss, \_, train\_state = Lux.Training.single\_train\_step!(AutoEnzyme(), custom\_loss,(g, g.x, g.y), train\_state)
> \`\`\`
> 
> I started seeing the cryptic errors again:
> \`\`\`julia
> ERROR: Constant memory is stored (or returned) to a differentiable variable.
> As a result, Enzyme cannot provably ensure correctness and throws this error.
> This might be due to the use of a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Runtime-Activity).
> If Enzyme should be able to prove this use non-differentable, open an issue!
> To work around this issue, either:
> a) rewrite this variable to not be conditionally active (fastest, but requires a code change), or
> b) set the Enzyme mode to turn on runtime activity (e.g. autodiff(set\_runtime\_activity(Reverse), ...) ). This will maintain correctness, but may slightly reduce performance.
> Mismatched activity for: store {} addrspace(10)\* %157, {} addrspace(10)\*\* %sret\_return.repack.repack, align 8, !dbg !234, !noalias !244 const val: %157 = call fastcc nonnull {} addrspace(10)\* @julia\_vcat\_89903({} addrspace(10)\* noundef nonnull align 8 dereferenceable(24) %15, {} addrspace(10)\* noundef nonnull align 8 dereferenceable(24) %59) #77, !dbg !228
> Type tree: {\[-1\]:Pointer, \[-1,0\]:Pointer, \[-1,0,-1\]:Integer, \[-1,8\]:Pointer, \[-1,8,0\]:Integer, \[-1,8,1\]:Integer, \[-1,8,2\]:Integer, \[-1,8,3\]:Integer, \[-1,8,4\]:Integer, \[-1,8,5\]:Integer, \[-1,8,6\]:Integer, \[-1,8,7\]:Integer, \[-1,8,8\]:Pointer, \[-1,8,8,-1\]:Integer, \[-1,16\]:Integer, \[-1,17\]:Integer, \[-1,18\]:Integer, \[-1,19\]:Integer, \[-1,20\]:Integer, \[-1,21\]:Integer, \[-1,22\]:Integer, \[-1,23\]:Integer}
> llvalue= %157 = call fastcc nonnull {} addrspace(10)\* @julia\_vcat\_89903({} addrspace(10)\* noundef nonnull align 8 dereferenceable(24) %15, {} addrspace(10)\* noundef nonnull align 8 dereferenceable(24) %59) #77, !dbg !228
> 
> Stacktrace:
> \[1\] add\_self\_loops
> @ /scratch/htc/amartine/julia/packages/GNNGraphs/Ldvz4/src/transform.jl:24
> 
> Stacktrace:
> \[1\] add\_self\_loops
> @ /scratch/htc/amartine/julia/packages/GNNGraphs/Ldvz4/src/transform.jl:24
> \[2\] #\_#10
> @ /scratch/htc/amartine/julia/packages/GNNLux/AHHiN/src/layers/conv.jl:137 \[inlined\]
> \[3\] GCNConv
> @ /scratch/htc/amartine/julia/packages/GNNLux/AHHiN/src/layers/conv.jl:131 \[inlined\]
> \[4\] GCNConv
> @ /scratch/htc/amartine/julia/packages/GNNLux/AHHiN/src/layers/conv.jl:128 \[inlined\]
> \[5\] GCN
> @ ./REPL\[21\]:3
> \[6\] custom\_loss
> @ ./REPL\[29\]:4
> \[7\] #4
> @ /scratch/htc/amartine/julia/packages/Lux/ptjU6/src/helpers/training.jl:257 \[inlined\]
> \[8\] augmented\_julia\_\_4\_89363\_inner\_33wrap
> @ /scratch/htc/amartine/julia/packages/Lux/ptjU6/src/helpers/training.jl:0
> \[9\] macro expansion
> @ /scratch/htc/amartine/julia/packages/Enzyme/rwbr4/src/compiler.jl:5610 \[inlined\]
> \[10\] enzyme\_call
> @ /scratch/htc/amartine/julia/packages/Enzyme/rwbr4/src/compiler.jl:5144 \[inlined\]
> \[11\] AugmentedForwardThunk
> @ /scratch/htc/amartine/julia/packages/Enzyme/rwbr4/src/compiler.jl:5083 \[inlined\]
> \[12\] autodiff
> @ /scratch/htc/amartine/julia/packages/Enzyme/rwbr4/src/Enzyme.jl:408 \[inlined\]
> \[13\] compute\_gradients\_impl
> @ /scratch/htc/amartine/julia/packages/Lux/ptjU6/ext/LuxEnzymeExt/training.jl:10 \[inlined\]
> \[14\] compute\_gradients
> @ /scratch/htc/amartine/julia/packages/Lux/ptjU6/src/helpers/training.jl:200 \[inlined\]
> \[15\] single\_train\_step\_impl!
> @ /scratch/htc/amartine/julia/packages/Lux/ptjU6/src/helpers/training.jl:327 \[inlined\]
> \[16\] #single\_train\_step!#6
> @ /scratch/htc/amartine/julia/packages/Lux/ptjU6/src/helpers/training.jl:292 \[inlined\]
> \[17\] single\_train\_step!(backend::AutoEnzyme{…}, obj\_fn::typeof(custom\_loss), data::Tuple{…}, ts::Lux.Training.TrainState{…})
> @ Lux.Training /scratch/htc/amartine/julia/packages/Lux/ptjU6/src/helpers/training.jl:285
> \[18\] train\_model!(gcn::GCN{…}, ps::@NamedTuple{…}, st::@NamedTuple{…}, g::GNNGraph{…})
> @ Main ./REPL\[35\]:4
> \[19\] top-level scope
> @ REPL\[36\]:1
> Some type information was truncated. Use \`show(err)\` to see complete types.
> \`\`\`
> 
> Now that I have come across #389 I realize why. I think it would be a good idea to include a warning on the documentation that Zygote.jl is not supported. Perhaps, better yet, catch it and let the user know before they get the scary-looking exception with bits of IR and so that they are less likely to spend a bunch of time wondering why it's just not working.

If you manage to get it to work, please let us know.  
You could also try Enzyme in combination with GNNLux.