Reverse rules for Enzyme

[kshyatt]

No description provided.

[codecov]

Codecov Report

❌ Patch coverage is 3.84615% with 150 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...ixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl	1.31%	150 Missing ⚠️

Files with missing lines	Coverage Δ
src/common/initialization.jl	100.00% <100.00%> (ø)
src/common/safemethods.jl	100.00% <100.00%> (ø)
...ixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl	1.31% <1.31%> (ø)

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

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

[github-actions]

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

[kshyatt]

OK, I think this is ready for review!

[kshyatt]

1.12 failures really seem unrelated/Enzyme internal, not sure what to do about it other than disable the Enzyme tests on 1.12 until it becomes more stable there.

[kshyatt]

OK, for unknown and likely extremely brain melting reasons, adding a repr(a) call specifically for ComplexF32 fixes all the problems with NaN we were seeing. It's upsetting and I'm not sure why, maybe inlining? Either way, while I investigate that, this should help unblock CI a little

[kshyatt]

Culprit seems to be inlining for ComplexF32 specifically, which is weird, and makes me unhappy, but at least we don't have to print to an unused string now

[kshyatt]

OK, made the caching of arg and ret a lot more consistent

[Jutho]

The end is near, I think. It is the first time I get all the way to the bottom of extension file 😄 .

One big question I have is that I still don't fully understand why we pass dret via the cache and how it differs from the dret argument (which is now no longer called like this) in the reverse function. Doesn't the Enzyme manual on defining reverse rules use the dret argument in the reverse function exactly like that:

https://enzyme.mit.edu/julia/stable/generated/custom_rule/#Defining-a-reverse-mode-rule

I didn't fully understand your previous response.

[kshyatt]

OK, so I think in the manual the only case in which we can use dret like that is when it is Active eg we return some float and we need access to its derivative. Otherwise, the 3rd argument in the reverse signature is a Type indicating what the annotation of the output was -- Const, Duplicated, BatchDuplicated, etc.

[kshyatt]

If the return is e.g. a Duplicated and freshly allocated we can put a reference to its derivative on the tape so that we can peek into it later for the pullback, which is what we're currently doing

[kshyatt]

Test fails are in lq_null/qr_null and the related orthnull tests. I'll take a look.

[kshyatt]

OK, I think we must always copy A.val for the qr_null or lq_null cases as those pullbacks do depend directly on A, NOT just on dA, N, and dN.

[Jutho]

Suggested change

	dret = (Diagonal(zero(diagview(DV′[1]))), zero(DV′[2]))
	shadow = EnzymeRules.needs_shadow(config) ? dret : nothing
	return EnzymeRules.AugmentedReturn(primal, shadow, (cache_DV, dret, ind))
	dret = if EnzymeRules.needs_shadow(config)
	(Diagonal(zero(diagview(DV′[1]))), zero(DV′[2]))
	else
	nothing
	end
	return EnzymeRules.AugmentedReturn(primal, dret, (cache_DV, dret, ind))

[Jutho]

While we should probably wait for the tests, I will approve prematurely 😄