Skip to content

Test: Replaced raw Mooncake with Flux implementation and checked Mooncake compatibility for more layers#642

Merged
CarloLucibello merged 3 commits intoJuliaGraphs:masterfrom
Parvm1102:mooncake-ad-testing
Mar 15, 2026
Merged

Test: Replaced raw Mooncake with Flux implementation and checked Mooncake compatibility for more layers#642
CarloLucibello merged 3 commits intoJuliaGraphs:masterfrom
Parvm1102:mooncake-ad-testing

Conversation

@Parvm1102
Copy link
Copy Markdown
Contributor

I had reported the friendly_tangents= true issue to Mooncake.jl and they have resolved it in Mooncake v0.5.17.
I have changed the implementation with the preferred Flux.withgradient().

Also I have also checked few more layers in basic.jl, pool.jl and temporal.jl.
Observations:

  • GlobalPool: PASS
  • GlobalAttentionPool: PASS
  • TGCNCell / TGCN: PASS
  • EvolveGCNO: PASS
  • GNNChain: Fail (MooncakeRuleCompilationError, complex chain with anonymous fns, Dropout, BatchNorm)
  • GConvLSTMCell / GConvLSTM: Fail (Uses ChebConv internally)
  • GConvGRUCell / GConvGRU: Fail (Uses ChebConv internally)
  • DCGRUCell / DCGRU: Fail (Uses DConv internally)
  • EvolveGCNOCell: Fail ("with state" call fails during Mooncake rule compilation)

…cake compatibility for more layers

Signed-off-by: Parvm1102 <parvmittal31757@gmail.com>
@Parvm1102 Parvm1102 force-pushed the mooncake-ad-testing branch from 2d8ad69 to b5130f0 Compare March 7, 2026 17:37
@Parvm1102
Copy link
Copy Markdown
Contributor Author

Parvm1102 commented Mar 7, 2026

@CarloLucibello could you please review this?

@Parvm1102
Copy link
Copy Markdown
Contributor Author

Interesting observation: TransformerConv and GATv2Conv were failing in the CI earlier, but now they pass with mooncake enabled.

@Parvm1102
Copy link
Copy Markdown
Contributor Author

@CarloLucibello Please review this when you have time

Comment thread GraphNeuralNetworks/test/test_module.jl Outdated
Comment thread GraphNeuralNetworks/test/test_module.jl Outdated
@CarloLucibello
Copy link
Copy Markdown
Member

GNNChain: Fail (MooncakeRuleCompilationError, complex chain with anonymous fns, Dropout, BatchNorm)

It would be good to have a MWE for this and report it upstream

@CarloLucibello CarloLucibello merged commit 3f134c9 into JuliaGraphs:master Mar 15, 2026
7 of 10 checks passed
g_mc_result = result.grad # Extract gradients tuple
y_mc, g_mc = Flux.withgradient(f -> loss(f, graph, xs...), Flux.AutoMooncake(), f)
@assert isapprox(y, y_mc; rtol, atol)
check_equal_leaves(g, g_mc_result; rtol, atol)
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
check_equal_leaves(g, g_mc_result; rtol, atol)
check_equal_leaves(g, g_mc; rtol, atol)

@CarloLucibello The g_mc_result should also be changed to g_mc, otherwise the tests are failing. Should I open a new pr fixing the mistake?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ops, sorry, yes please fix

Parvm1102 added a commit to Parvm1102/GraphNeuralNetworks.jl that referenced this pull request Mar 15, 2026
Signed-off-by: Parvm1102 <parvmittal31757@gmail.com>
@Parvm1102 Parvm1102 mentioned this pull request Mar 15, 2026
CarloLucibello pushed a commit that referenced this pull request Mar 16, 2026
Signed-off-by: Parvm1102 <parvmittal31757@gmail.com>
@Parvm1102
Copy link
Copy Markdown
Contributor Author

It would be good to have a MWE for this and report it upstream

Yeah I wil do it soon!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants