Pullback rules for truncation methods

[Jutho]

A first attempt at generalizing the eig, eigh and svd pullback rules so that they can also be used in the case of truncated methods with TruncatedAlgorithm, i.e. where first the full decomposition is computed. This was already supported for svd, but with the implicit assumption that the singular values/vectors kept were always the first ones.

This clearly needs additional tests, especially for the eig and eigh methods where such a sorted truncation order is definitely not guaranteed.

Alternatively, I can also write an implementation that does not use the full decomposition, and solves the problem in terms of Sylvester equations with the original matrix. That might actually more generic. In particular, I am also thinking about randomized SVD, where one would not have the full decomposition available (and for which the current rules would thus be wrong or failing).

[codecov]

Codecov Report

❌ Patch coverage is 94.47005% with 12 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/pullbacks/eigh.jl	92.72%	4 Missing ⚠️
src/pullbacks/svd.jl	95.71%	3 Missing ⚠️
ext/MatrixAlgebraKitChainRulesCoreExt.jl	93.33%	2 Missing ⚠️
src/pullbacks/eig.jl	96.42%	2 Missing ⚠️
src/pullbacks/lq.jl	75.00%	1 Missing ⚠️

Files with missing lines	Coverage Δ
src/MatrixAlgebraKit.jl	100.00% <ø> (ø)
src/pullbacks/polar.jl	100.00% <100.00%> (ø)
src/pullbacks/qr.jl	96.22% <ø> (ø)
src/pullbacks/lq.jl	96.29% <75.00%> (ø)
ext/MatrixAlgebraKitChainRulesCoreExt.jl	83.73% <93.33%> (+0.56%)	⬆️
src/pullbacks/eig.jl	95.94% <96.42%> (-0.36%)	⬇️
src/pullbacks/svd.jl	96.26% <95.71%> (-1.66%)	⬇️
src/pullbacks/eigh.jl	91.04% <92.72%> (-3.96%)	⬇️

🚀 New features to boost your workflow:

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

[Jutho]

As @kshyatt is working on package extensions for supporting Enzyme.jl and Mooncake.jl, it might be good to revisit the interface we have currently chosen to write the pullbacks within the MatrixAlgebraKit module itself. The current interface consists of methods having as arguments:

• the adjoint of the input matrix (probably matrices for generalized eigenvalue problems)
• the output data as a single (tuple) argument
• the adjoints of the output as a single (tuple) argument
• additional arguments and keyword arguments to control behavior of the pullback computation
  Furthermore, there is a single iszerotangent function that can be overloaded by package extensions.

While this has worked well so far in combination with ChainRules.jl & Zygote.jl, it is apparently not the best fit for the other AD packages. As far as I understood, these other AD packages encode the absence/zeroness of certain adjoints differently, and this conflicts with the fact that the different incoming adjoints (those associated with the output variables) in our interface are assumed to be packed in a tuple.

Hence, I want to put this interface back on the table before continuing with this PR. Input from @kshyatt will be essential.

[Jutho]

One other problem while implementing the alternative methods for the truncated decompositions (that do not require the full decomposition) is that they do need the input matrix A. This argument is of course available, it was just not necessary for the other decompositions.

[lkdvos]

Definitely happy to rethink the interface for the pullbacks. I don't think we specifically had many things in mind when we designed this, so restrategizing now that Mooncake and Enzyme are on the horizon seems like the right thing to do.

[Jutho]

I am wondering whether modifying the incoming adjoints in place is allowed by all the AD backends. That is why I had refrained from doing so. Maybe @kshyatt can comment on this?

[lkdvos]

You are probably right, and I definitely don't want to risk that for saving a single allocation. I just hadn't thought about that.