Mooncake 0.5.25 compat by sunxd3 · Pull Request #988 · JuliaDiff/DifferentiationInterface.jl

sunxd3 · 2026-04-04T14:47:35Z

An attempt at addressing #986.

Feel free to make any edits or take over!

codecov · 2026-04-04T14:53:27Z

Codecov Report

❌ Patch coverage is 90.47619% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 97.28%. Comparing base (a5ecbe0) to head (ce72baf).

Files with missing lines	Patch %	Lines
...e/ext/DifferentiationInterfaceMooncakeExt/utils.jl	81.81%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #988      +/-   ##
==========================================
- Coverage   98.21%   97.28%   -0.94%     
==========================================
  Files         135      131       -4     
  Lines        8000     7984      -16     
==========================================
- Hits         7857     7767      -90     
- Misses        143      217      +74

Flag	Coverage Δ
DI	`97.86% <90.47%> (-1.12%)`	⬇️
DIT	`95.76% <ø> (-0.47%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

sunxd3 · 2026-04-04T20:31:01Z

I think I understand the CI error, there is something I need to patch on the Mooncake side, will come back to this

sunxd3 · 2026-04-09T06:50:22Z

chalk-lab/Mooncake.jl#1129 should unblock this PR, we'll release it as soon as it's available

gdalle · 2026-04-10T06:37:18Z

Thank you for taking a crack at this! I'll wait until the tests pass before reviewing if that's okay

sunxd3 · 2026-04-10T07:05:22Z

totally fine!

Mooncake returns raw Tangent objects instead of friendly arrays for StaticArrays on Julia 1.11. This is an upstream bug — skip the test until it is fixed.

On Julia 1.11, Mooncake may return raw Tangent objects instead of friendly arrays for StaticArrays even with friendly_tangents=true. Add _maybe_to_primal dispatch as a safety net that converts leaked Tangent objects to primal-shaped values, no-op otherwise.

Also convert leaked Mooncake.MutableTangent (e.g. MVector tangents) and apply _maybe_to_primal in forward mode (pushforward) paths.

sunxd3 · 2026-04-10T13:40:32Z

@gdalle the Mooncake CIs are passing (it probably requires Mooncake v0.5.26).

Could you take over? I also won't be offended if you want to start a new PR.

gdalle · 2026-04-10T14:24:27Z

I'll take a look when I can, thanks a bunch! DI's tests are failing on main too because of Mooncake's breaking release so this is a priority for me. Do you know why coverage is not complete?

sunxd3 · 2026-04-10T15:11:17Z

Thanks a lot.

Do you know why coverage is not complete?

I am not certain. A bad guess is that some code changes are more defensive than necessary. Sorry!

gdalle

Thank you for trying to fix what others broke! I added a few remarks to understand the task a bit better, I'll wait for your answers

gdalle · 2026-04-16T12:22:37Z

+@inline maybe_getfield(mod, name::Symbol) =
+    isdefined(mod, name) ? getfield(mod, name) : nothing
+
+const mooncake_tangent_to_friendly = maybe_getfield(
+    Mooncake, Symbol("tangent_to_friendly!!")
+)
+const mooncake_friendly_tangent_cache = maybe_getfield(Mooncake, :FriendlyTangentCache)
+const mooncake_as_primal = maybe_getfield(Mooncake, :AsPrimal)
+const mooncake_no_cache = maybe_getfield(Mooncake, :NoCache)


This doesn't seem to be very robust? I'd rather impose a lower bound for Mooncake at v0.5.25 in Project.toml (that way we're sure we can use all of these symbols)

gdalle · 2026-04-16T12:23:19Z

        )
        y = first(y_and_dy)
-        dy = _copy_output(last(y_and_dy))
+        dy = _maybe_to_primal(last(y_and_dy), y)


Why do we need to ensure that primal conversion happens here? If friendly_tangents is set to true, won't Mooncake's pushforward and pullback already return a primal-like object?

gdalle · 2026-04-16T12:25:26Z

+    backend = AutoMooncake(; config = Mooncake.Config(; friendly_tangents = true))
+    inputs = (
+        Symmetric([2.0 1.0; 1.0 3.0]),
+        Hermitian(ComplexF64[2 1 + im; 1 - im 3]),


Do you know which convention Mooncake uses for gradients of functions with complex inputs and real outputs? There are two possible choices, see e.g. https://arxiv.org/abs/2409.06752

See: chalk-lab/Mooncake.jl#773

gdalle · 2026-04-16T12:26:12Z

+    @testset "$(typeof(x))" for x in inputs
+        grad = gradient(f, backend, x)
+        y, grad2 = value_and_gradient(f, backend, x)
+        pb = only(pullback(identity, backend, x, (x,)))


This is not a strong enough test, the function is too simple

gdalle · 2026-04-16T12:26:25Z

+        @test grad isa Matrix
+        @test grad2 isa Matrix
+        @test pb isa Matrix
+        @test grad == grad2


grad and grad2 are never compared against the ground truth

gdalle · 2026-04-16T12:27:16Z

+            !isnothing(mooncake_as_primal) &&
+            !isnothing(mooncake_no_cache)
+        dest = mooncake_friendly_tangent_cache{mooncake_as_primal}(_copy_output(x))
+        cache = isbitstype(typeof(x)) ? mooncake_no_cache() : IdDict{Any, Any}()


Why is a type-unstable dictionary needed here?
Does this make every tangent-to-primal conversion outside of bitstypes slow?

Regardless of why, we may want to allocate this dictionary in the preparation phase

yup this will regress performance on hot paths. This should be allocated once during the prepare phase and stored in the extras cache.

gdalle · 2026-04-16T12:29:40Z

@AstitvaAggarwal @Technici4n could you maybe take a look too?

AstitvaAggarwal

also we might want to keep track of future possible tangent_types: _maybe_to_primal(x, _) = x will silently pass through any tangent type not yet accounted for (e.g. a future Mooncake.SparseTangent), making failures invisible.

AstitvaAggarwal · 2026-04-17T09:03:42Z

+            !isnothing(mooncake_as_primal) &&
+            !isnothing(mooncake_no_cache)
+        dest = mooncake_friendly_tangent_cache{mooncake_as_primal}(_copy_output(x))
+        cache = isbitstype(typeof(x)) ? mooncake_no_cache() : IdDict{Any, Any}()


yup this will regress performance on hot paths. This should be allocated once during the prepare phase and stored in the extras cache.

Replace the broad type-unstable _maybe_to_primal shim (with its IdDict fallback and defensive symbol lookups) with a narrow _to_friendly_value hook in the Mooncake extension and a new triple-extension activated by [Mooncake, StaticArrays]. Mooncake's friendly_tangents framework (chalk-lab/Mooncake.jl#1103) deliberately defaults to :as_raw for SArray/MArray primals; bridge that gap at the DI boundary where the primal type is statically known. Reconstruction via typeof(x)(Mooncake.val(t.fields.data)) is unambiguous because the data::NTuple field maps 1:1 to logical positions, with no aliasing. Restrict the dispatch to scalar float / complex-float eltypes so static arrays with non-float eltype keep using Mooncake's element-wise AbstractArray recursion. Route zero_tangent_or_primal through the same hook so the prep-time dy_backup buffer is primal-shaped (otherwise copyto!(::MutableTangent, ::MVector) fails at twoarg.jl:54). Restore nomatrix(static_scenarios()) on the friendly_tangents=true backends.

gdalle · 2026-04-26T18:14:56Z

@sunxd3 can you explain why type piracy is needed here? Sounds like there should be a StaticArraysExt in Mooncake (but the maintainer doesn't like it) or the other way around. I will not introduce piracy in a package depended upon by the entire ecosystem.

yolhan83 · 2026-04-26T19:49:33Z

would be weird but one could do a thin DI wrapper DISArray of a SArray which would just wrap a SArray everytime it needs to go into a Mooncake function, llvm should understand this is unneeded and remove it at compile time (since its all static) but then, there is the question of do we need any kind of interface on those, and I hope not otherwise its just too big of a job.

sunxd3 added 3 commits April 4, 2026 15:37

fix Mooncake friendly_tangents compatibility

75d3c05

format Mooncake fix

f420234

remove Mooncake test formatting churn

8235e6c

sunxd3 requested a review from gdalle as a code owner April 4, 2026 14:47

gdalle marked this pull request as draft April 5, 2026 08:36

sunxd3 added 5 commits April 10, 2026 08:12

ci: retrigger after Mooncake v0.5.26 release

a478578

style: apply Runic formatting to utils.jl

65997c4

test: skip friendly_tangents static_scenarios on Julia 1.11

106f50f

Mooncake returns raw Tangent objects instead of friendly arrays for StaticArrays on Julia 1.11. This is an upstream bug — skip the test until it is fixed.

fix: handle MutableTangent and forward mode tangent leaks

ce72baf

Also convert leaked Mooncake.MutableTangent (e.g. MVector tangents) and apply _maybe_to_primal in forward mode (pushforward) paths.

This was referenced Apr 11, 2026

Mooncake ext: unwrap struct tangents back to primal type #989

Closed

Changes in Mooncake 0.5.25 that may affect DI users #986

Open

gdalle requested changes Apr 16, 2026

View reviewed changes

AstitvaAggarwal reviewed Apr 17, 2026

View reviewed changes

This was referenced Apr 20, 2026

chore(deps): bump julia-actions/setup-julia from 2 to 3 #994

Merged

chore: upper-bound Mooncake version (coherence with General) #995

Merged

ChrisRackauckas-Claude mentioned this pull request Apr 22, 2026

Unpin / widen test/ad Project.toml compat to resolve Mooncake ↔ SciMLSensitivity conflict SciML/OrdinaryDiffEq.jl#3503

Closed

3 tasks

Conversation

sunxd3 commented Apr 4, 2026

Uh oh!

codecov Bot commented Apr 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

sunxd3 commented Apr 4, 2026

Uh oh!

sunxd3 commented Apr 9, 2026

Uh oh!

gdalle commented Apr 10, 2026

Uh oh!

sunxd3 commented Apr 10, 2026

Uh oh!

sunxd3 commented Apr 10, 2026

Uh oh!

gdalle commented Apr 10, 2026

Uh oh!

sunxd3 commented Apr 10, 2026

Uh oh!

gdalle left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gdalle commented Apr 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

AstitvaAggarwal left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gdalle commented Apr 26, 2026

Uh oh!

yolhan83 commented Apr 26, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

codecov Bot commented Apr 4, 2026 •

edited

Loading

gdalle commented Apr 16, 2026 •

edited

Loading