Check that the correct AD backend is being used

[mhauru]

Introduces tests that check that the ADType we specify is actually the one being used. This alleviates my worry that we might mess up passing the ADType around and accidentally fall back on the default.

I've added tests for calling sample and maximum_likelihood/maximum_a_posteriori. Any other user-facing functions we should check to make sure they, too, use their adtype arguments correctly?

Closes #2235

[mhauru]

@torfjelde pointed out in https://github.com/TuringLang/Turing.jl/pull/2221/files#r1608405934 that there's a similarity between this and DebugContext, in that they both call their child context as usual and just add an extra call before/after. Thus we could rather implement a more general CallbackContext, and unify the implementation of DebugContext and ADTypeCheckContext under it.

I think this is a great idea, but I chose not to do it in this PR because

1. Would need to give a good think as to what powers the callbacks should have. Are they e.g. allowed to modify varinfo?
2. Would want to check whether any other contexts we currently have fit this CallbackContext format, and bring those in as well.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 85.77%. Comparing base (29a1342) to head (15701c2).
Report is 3 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2291      +/-   ##
==========================================
- Coverage   85.82%   85.77%   -0.06%     
==========================================
  Files          24       24              
  Lines        1623     1617       -6     
==========================================
- Hits         1393     1387       -6     
  Misses        230      230              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[sunxd3]

The code looks good to me.

I don't think comparing the types of values in varinfo is particularly elegant, but also I don't see other ways to do it with Contexts.

[sunxd3]

Maybe @yebai, @torfjelde, or @devmotion can give a quick scan?

[coveralls]

Pull Request Test Coverage Report for Build 10268798219

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

• For more information on this, see Tracking coverage changes with pull request builds.
• To avoid this issue with future PRs, see these Recommended CI Configurations.
• For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

• 0 of 0 changed or added relevant lines in 0 files are covered.
• 2 unchanged lines in 1 file lost coverage.
• Overall coverage decreased (-0.05%) to 85.989%

---

Files with Coverage Reduction	New Missed Lines	%
src/mcmc/abstractmcmc.jl	2	94.55%

Totals
Change from base Build 9953741319:	-0.05%
Covered Lines:	1387
Relevant Lines:	1613

---

💛 - Coveralls

[mhauru]

I don't think comparing the types of values in varinfo is particularly elegant, but also I don't see other ways to do it with Contexts.

Agree on lack of elegance, better proposals welcome.

[torfjelde]

I like the idea very much:) Seems like something we should definitively test. And agree that we shouldn't do the CallbackContext in this PR.

I left a comment on a potentially way we can check Zygote.jl, Tapir.jl, and Enzyme.jl too.

[torfjelde]

Yeah there's no way you can do this here unfortunately 😕

Possibly crazy idea, but it might make sense to specifically overload the adjoint computation for a given backend which doesn't use types, e.g.

Zygote.@adjoint function check_adtype(...)
    # Should only be hit if we're using Zygote.jl.
    ...
end

? Could do the same with Enzyme.jl and Tapir.jl, both of which would suffer from the same issue

[mhauru]

I like this, and implemented it. Turns out there's Zygote.hook for exactly these sorts of things.

I thought a bit about switching to doing all checks like this, and giving up on the element type approach, but that seemed like it would get more complicated than is worth at this point. We now catch all cases except if one uses ForwardDiff when Zygote is expected, which seems sufficient to me.

[torfjelde]

Very, very nice @mhauru :)

[mhauru]

The only CI failure is a tiny, trivial loss of code coverage. @torfjelde, are you happy to have this merged as is?

[torfjelde]

Very, very nice @mhauru :)