Add Mooncake and Enzyme tests for Diagonal by kshyatt · Pull Request #179 · QuantumKitHub/MatrixAlgebraKit.jl

kshyatt · 2026-02-27T10:24:48Z

No description provided.

codecov · 2026-02-27T10:41:21Z

Codecov Report

❌ Patch coverage is 42.71845% with 59 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...ixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl	0.00%	43 Missing ⚠️
...gebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl	73.33%	16 Missing ⚠️

Files with missing lines	Coverage Δ
...gebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl	`64.20% <73.33%> (-0.11%)`	⬇️
...ixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl	`1.12% <0.00%> (-0.17%)`	⬇️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

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

kshyatt · 2026-03-26T12:44:53Z

Went a little crazy with this one after figuring out some of the rules needed modification for the Diagonal case where the output is by default the same as the input

github-actions · 2026-03-26T12:45:14Z

Your PR no longer requires formatting changes. Thank you for your contribution!

kshyatt · 2026-03-27T17:04:11Z

This was a fun one, as a bunch of rules had to be modified to account for the case where one of the arguments === the input. Seems to all be working now, at least!

lkdvos

Good catch, that does indeed sound like an absolute misery to debug/find and solve!

I left some small suggestions around code organization, mostly to reduce the amount of code and/or improve readability, but otherwise this looks great (I absolutely understand that the complexity might just be unavoidable, again, impressive catch on this!)

I'll go through the remainder afterwards if you like, but I will already leave this here to not get lost in a wall of questions :)

lkdvos · 2026-03-30T13:17:24Z

ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

+                if A_is_arg
+                    ΔA = make_zero(A.dval)
+                    $pb(ΔA, Aval, argval, darg)
+                    A.dval .= ΔA


Suggested change

A.dval .= ΔA

A.dval .+= ΔA

I'm guessing this has to be accumulating, but I'm not entirely sure I follow the logic of this branch, why can't we immediately in-place accumulate?

We can't do it inplace because you'll modify dA which is the same as darg, I think. I think we should not do a + because A is overwritten in the primal step (but this is confusing, I'm open to being convinced otherwise)

I think I'm following now.
Would it help if instead of allocating a new dA and then doing the copy afterwards, we instead copy the darg if we need to?
Also, does that mean that we could in principle be checking that dA is zero, since nothing is allowed to depend on the value of A after the decomposition?
Or would that be too strict of a check in the case where A is secretly not overwritten by the specific implementation of the decomposition.

Let me think about this a bit... btw this whole hornet's nest would be much easier to handle if we finally implemented the qr_full!! stuff we had discussed at one point 😛

I thought Jutho wasn't a big fan of that 🤔

Then the hornets are sticking around!

ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

lkdvos · 2026-03-30T13:29:44Z

ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

            # use A (so that whoever does this is forced to handle caching A
            # appropriately here)
            Aval = nothing
+            A_is_arg = !isa(A, Const) && TA <: Diagonal && diagview(A.dval) === D.dval


Code-organization-wise, I am not a huge fan of this hardcoded special case for different types.
Why do you need the specific specialization here? Could it not just be A.dval === D.dval?

That being said, we should probably just add an implementation of isalias(A, arg).
Maybe something like Base.mightalias could be a more generic solution, or even a fallback definition, but that function is technically internal, and similar problems hold for checking Base.dataids.
I might actually be okay with depending on that though, that has been quite stable and does seem to me like a good way of going at this. ( and we already secretly use this in TensorOperations! )

For the _vals methods we can't do A.dval === D.dval because D is a Vector, and A is a Diagonal, whose diag field points to D

Ah you are right, my bad! I think then I would even more think an isalias(A, arg) = Base.mightalias(A, arg) is the right abstraction (I would leave in the hook since Base.mightalias is internal, and this allows us to overload without piracy)

Yeah, that sounds good to me. We could do it as part of this PR or separately?

There are a lot of goofy special cases running around here haha

It might be reasonable to include it in this PR, as this does seem to be the primary change that is needed to solve the issue? Happy to defer it if you prefer though.

TBH this PR is complex enough, I might do the isalias as a followup?

lkdvos · 2026-03-30T13:35:36Z

ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

+                if !(A === arg1 || A === arg2)
+                    copy!(A, Ac)
+                    $pb(dA, A, (arg1, arg2), (darg1, darg2))
+                else
+                    ΔA = zero(A)
+                    $pb(ΔA, A, (arg1, arg2), (darg1, darg2))
+                    dA .= ΔA
+                end


Suggested change

if !(A === arg1 || A === arg2)

copy!(A, Ac)

$pb(dA, A, (arg1, arg2), (darg1, darg2))

else

ΔA = zero(A)

$pb(ΔA, A, (arg1, arg2), (darg1, darg2))

dA .= ΔA

end

# compute pullback of A and restore input

# order is important - A might alias `args`!

$pb(dA, Ac, (arg1, arg2), (darg1, darg2))

copy!(A, Ac)

# compute pullbacks and restore other inputs if they don't alias

if A !== arg1

copy!(arg1, arg1c)

zero!(darg1)

end

if A !== arg2

copy!(arg2, arg2c)

zero!(darg2)

end

Again mostly a readability suggestion. (Note that I might again be missing something about the use of dA vs Delta A here)

See comment above about Delta A

suggestion for the bottom part still stands though :)

NB that if A !== arg, we MUST copy Ac -> A before the pullback

Pushed a small cleanup here at least...

kshyatt · 2026-03-31T12:16:53Z

ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

-                elseif A_is_arg2
-                    make_zero!(arg.dval[1])
-                end
+                A.dval === arg.dval[1] || make_zero!(arg.dval[1])


I think this is actually not correct...

Or at least seems to be causing test failures now?

kshyatt · 2026-03-31T14:16:31Z

CUDA fail is unrelated, it's the CUSOLVER memory thing again

Fix stupid typos Diagonal + eig(h) working SVD working for diag LQ/QR working Polar and orthnull Some Mooncake progress Formatter Working Mooncake eig Remove comment No diag for orthnull Diag tests working for Mooncake Formatter Update ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl Co-authored-by: Lukas Devos <ldevos98@gmail.com> A bit of Enzyme cleanup A little Mooncake cleanup

kshyatt requested a review from lkdvos February 27, 2026 10:24

kshyatt force-pushed the ksh/diag branch from 5e9ee7f to 415f3f4 Compare March 23, 2026 11:44

kshyatt changed the title ~~Add Mooncake tests for Diagonal~~ Add Mooncake and Enzyme tests for Diagonal Mar 26, 2026

kshyatt force-pushed the ksh/diag branch from 77c632d to cb3a53c Compare March 26, 2026 16:20

kshyatt marked this pull request as ready for review March 27, 2026 16:59

lkdvos reviewed Mar 30, 2026

View reviewed changes

kshyatt force-pushed the ksh/diag branch from 785f9fa to 187a4e6 Compare March 31, 2026 09:50

kshyatt commented Mar 31, 2026

View reviewed changes

kshyatt force-pushed the ksh/diag branch from 187a4e6 to 5d4c18e Compare March 31, 2026 13:36

kshyatt force-pushed the ksh/diag branch from f854821 to 65bb3c2 Compare April 2, 2026 13:12

-                if !(A === arg1 || A === arg2)
-                    copy!(A, Ac)
-                    $pb(dA, A, (arg1, arg2), (darg1, darg2))
-                else
-                    ΔA = zero(A)
-                    $pb(ΔA, A, (arg1, arg2), (darg1, darg2))
-                    dA .= ΔA
-                end
+                # compute pullback of A and restore input
+                # order is important - A might alias `args`!
+                $pb(dA, Ac, (arg1, arg2), (darg1, darg2))
+                copy!(A, Ac)
+                # compute pullbacks and restore other inputs if they don't alias
+                if A !== arg1
+                    copy!(arg1, arg1c)
+                    zero!(darg1)
+                end
+                if A !== arg2
+                    copy!(arg2, arg2c)
+                    zero!(darg2)
+                end

Conversation

kshyatt commented Feb 27, 2026

Uh oh!

codecov bot commented Feb 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

kshyatt commented Mar 26, 2026

Uh oh!

github-actions bot commented Mar 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kshyatt commented Mar 27, 2026

Uh oh!

lkdvos 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!

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!

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!

kshyatt commented Mar 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov bot commented Feb 27, 2026 •

edited

Loading

github-actions bot commented Mar 26, 2026 •

edited

Loading