Make implicit diff use the symmetry properties of C4v

[kshyatt]

I'm not sure how much of a performance benefit this is yet, but I think it should save us some memory at large chi. The idea is to exploit the 4-fold symmetry so that we only have to do the fixed point diff operations on 2 objects, rather than 8. I defined a new tangent helper type for this and some related VI helpers.

[Yue-Zhengyuan]

Judging by the look of C4vEnvTangent, maybe just introduce a dedicated C4vCTMEGEnv (which we've always wanted to do) ?

[kshyatt]

Yeah I was a bit torn on this. I didn't follow the original C4v PR closely enough to be certain that was what the consensus was

[pbrehmer]

Judging by the look of C4vEnvTangent, maybe just introduce a dedicated C4vCTMEGEnv (which we've always wanted to do) ?

Indeed, I think there was a consensus that we want to have specialized environment types at some point, including a C4vCTMRGEnv, the last time we discussed this. Although at that time we didn't implement it right away because it might entail a larger rewrite and some general design choices we hadn't settled on yet. Personally, I would also prefer adding a specialized C4vCTMRGEnv type to take care of this issue since that would be cleaner in the long run and could also generalize to other specialized environment types.

[kshyatt]

In that case I will close this and start on that and we can debate the interface and design stuff on a separate PR :)

[leburgel]

I think if we go through the trouble of adding a VI-compatible struct, it's better to go all the way and set up a dedicated C4vCTMEGEnv. If we don't want to break things but still want a smaller linear problem here, we could wrap the variables in a plain tuple and add some input and output hooks for the iteration VJP. The only issue is that VI doesn't handle tuples with mixed-scalartype entries, so there would have to be some hacky promotion.

Whatever we end up doing, I would prefer to add the hooks in a way that's a no-op for the non-C4v case, so we don't have to actually copy the rule body for different cases.

[codecov]

Codecov Report

❌ Patch coverage is 96.29630% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...lgorithms/optimization/implicit_differentiation.jl	96.29%	2 Missing ⚠️

Files with missing lines	Coverage Δ
...lgorithms/optimization/implicit_differentiation.jl	91.72% <96.29%> (+1.00%)	⬆️

... 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.