Support Mixing AD Frameworks for LogDensityProblems and the objective

[Red-Portal]

This PR enables mixing AD frameworks for the target LogDensityProblem and differentiating through the variational objective. This is enabled by defining a custom rrule through ChainRulesCore. As such, we are restricted to AD frameworks that support importing rrules. However, we still support directly differentiating through LogDensityProblems.logdensity if the target LogDensityProblem only has zeroth-order capability.

This PR depends on #179

[yebai]

Very helpful clarification, @Red-Portal. I understand your motivation now.

Therefore, the other option, which is taken by this PR, is to AD through $\mathrm{logdensity}(T_{\lambda}(\epsilon))$ as before, but provide an rrule (defined via ChainRulesCore) that calls logdensity_and_gradient. This enables the efficient JVP calculations, which properly uses the LogDensityProblems interface.

IIRC, logdensity_and_gradient is a LogDensityProblem API. So it is still okay to implement logdensity_and_gradient via DI, instead of bringing back the LogDensityProblemAD dependency. Correct?

[Red-Portal]

IIRC, logdensity_and_gradient is a LogDensityProblem API. So it is still okay to implement logdensity_and_gradient via DI, instead of bringing back the LogDensityProblemAD dependency. Correct?

Yes, the new dependency is only ChainRulesCore. I am not sure which LogDensityProblemAD dependency you are mentioning. The main Project.toml does not rely on it, only bench/Project.toml, docs/Project.toml and test/Project.toml do; for calculating the gradients of the test case/benchmark problem. Are you suggesting not to use it for tests/, docs/, and bench?

[Red-Portal]

@yebai Gentle nudge

[yebai]

only bench/Project.toml, docs/Project.toml and test/Project.toml do; for calculating the gradients of the test case/benchmark problem. Are you suggesting not to use it for tests/, docs/, and bench?

Yes, preferably we can remove dependencies on LogDensityProblemsAD in the tests as well.

I'd also prefer splitting this PR into two ones as @penelopeysm suggested in Slack.

I understand that it is a bit demanding, but it is very useful to gradually improve the quality of AdvancedVI.

[Red-Portal]

Hi @yebai , thanks for the comments.

Yes, preferably we can remove dependencies on LogDensityProblemsAD in the tests as well.

That would be doable by calling ForwardDiff directly.

I'd also prefer splitting this PR into two ones as @penelopeysm suggested in Slack.

In this PR's case, there are no file renaming/moving. I could split the doc updates to a different PR if that's what you meant.

[yebai]

That would be doable by calling ForwardDiff directly.

Can you use DI.gradient(...) or DI.hessian(...)?

In this PR's case, there are no file renaming/moving. I could split the doc updates to a different PR if that's what you meant.

@penelopeysm's wrote on slack:

I think it's mostly the commits that shift files around, like the first three, that would be easier to handle if it was in a separate PR.
Otherwise you get diffs that kind of look like +500 lines here and -499 lines there, and it can be hard to see what was really changed