Add code for AbstractMCMC Callbacks support

[shravanngoswamii]

No description provided.

[github-actions]

Turing.jl documentation for PR #2755 is available at:
https://TuringLang.github.io/Turing.jl/previews/PR2755/

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 85.17%. Comparing base (e487ed0) to head (58a4496).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2755      +/-   ##
==========================================
+ Coverage   85.12%   85.17%   +0.05%     
==========================================
  Files          20       20              
  Lines        1304     1309       +5     
==========================================
+ Hits         1110     1115       +5     
  Misses        194      194              

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

🚀 New features to boost your workflow:

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

[penelopeysm]

As far as I can tell, the main reason for all of this code is to satisfy names_and_values, which is only used in the TensorBoard callback.

Firstly, I think the AbstractMCMC part should be changed. names_and_values should be moved into the TensorBoard extension. It's not used anywhere else in AbstractMCMC/src, and it's not meant to be overloaded. On the bright side, you did declare it internal so you can do this without calling it a breaking change in semver.

Once that's done, we "only" need to overload the following methods:

• getparams and getstats -- okay (modulo my comment below about the return type of getparams)
• hyperparams_metric and _hyperparams_impl -- not so okay. These functions are not exported, and are not obviously part of the AbstractMCMC API. The latter has a comment above it saying "internal". That means that if AbstractMCMC changes them in a patch or minor release, which could conceivably be semver compliant from the AbstractMCMC side, then Turing code will suddenly stop working. This is actually quite a common occurrence: just a week ago I had to fix this case where we overloaded a DynamicPPL internal function #2748.

Before this can be merged into Turing, I would at least like to see the hyperparams_... functions in AbstractMCMC be made part of a public API.

I'm less fussed about this, but I also think that it should also be made really explicit, by naming, that their existence is solely for the TensorBoard logging. That is, in src you can do:

function tb_hyperparam_metric end

and then the default definition can live in AbstractMCMCTensorBoardExt.

[yebai]

Thanks @shravanngoswamii — I’ve added a few refactoring comments above. The functionality is already in place, but it could benefit from some minor cleanup, which I agree with Penny on.

[shravanngoswamii]

CI failure is not relevent to this PR.

[yebai]

Thanks, @shravanngoswamii -- I'm happy with this PR!

[penelopeysm]

Thanks @shravanngoswamii -- this is looking nice and clean now!

There is one point that I think we should discuss, which is the meaning of the 'params' and 'stats'. In particular, for samplers which operate in linked space, what do we return? e.g. does getparams return a vector where x is always between 0 or 1? (if this is confusing, there's an explanation here https://turinglang.org/docs/developers/transforms/distributions/)

@model function f()
    x ~ Beta(2, 2)
end

Currently, the implementation of getparams and getstats returns whatever the sampler stores internally. That means that some samplers will have unlinked values (e.g. default MH, which samples statically from the prior in unlinked space) and some samplers will have linked values (e.g. NUTS). The same is true of stats -- currently we return the sampler's notion of lp which may or may not be linked.

I don't mind as long as we are consistent, for example, we could say that getstats and getparams for Turing samplers always return whatever is internal to the sampler -- but I do think that it is much better from a user perspective to get the unlinked values since the internal values have no external meaning. We used to return internal log probs in the MCMCChains object, but we switched it over to returning 'true' unlinked lps recently, and I think that was a good thing.

[shravanngoswamii]

Thank you @penelopeysm.

getparams(state) returns get_varinfo(state)[:]:

• HMC/NUTS/SGHMC/SGLD: Linked (unconstrained) -- these samplers explicitly link the VarInfo
• MH() (default/static proposals): Unlinked -- should_link returns false
• MH(covariance_matrix) or explicit RandomWalkProposal: Linked
• ESS/PG: Unlinked

getstats(state).lp is also inconsistent:

• HMC: sum(getlogp(vi)) = logprior + logjac + loglikelihood
• MH: logjoint_internal = logprior + loglikelihood - logjac

However, ParamsWithStats (the transition) already does the right thing -- DynamicPPL.ParamsWithStats re-evaluates the model via ValuesAsInModelAccumulator and returns unlinked parameter values with proper logjoint (no Jacobian artifacts).

For callback-based monitoring, users interact with ParamsWithStats, not the raw state. So the user-facing behavior is correct.

Maybe, we can document that getparams/getstats return internal sampler state (which may be in transformed space), while ParamsWithStats provides semantically correct values. This matches the intent that these are low-level accessors vs. high-level transition containers.

Happy to add this documentation if that sounds reasonable.

[shravanngoswamii]

Tests will pass after TuringLang/AbstractMCMC.jl#191

> julia --project test\mcmc\callbacks.jl          
┌ Info: Found initial step size
└   ϵ = 1.6500000000000001
┌ Info: Found initial step size
└   ϵ = 1.6500000000000001
Test Summary:                    | Pass  Total   Time
AbstractMCMC Callbacks Interface |   80     80  38.3s

[penelopeysm]

However, ParamsWithStats (the transition) already does the right thing -- DynamicPPL.ParamsWithStats re-evaluates the model via ValuesAsInModelAccumulator and returns unlinked parameter values with proper logjoint (no Jacobian artifacts). For callback-based monitoring, users interact with ParamsWithStats, not the raw state. So the user-facing behavior is correct.

Does it? Don't all the methods getstats(...), etc., act on the state, not the transition? And that's what the callback framework uses?

[shravanngoswamii]

Does it? Don't all the methods getstats(...), etc., act on the state, not the transition? And that's what the callback framework uses?

I think the callback behavior is correct. The TensorBoard callback in AbstractMCMC calls ParamsWithStats(model, sampler, transition, state) and Turing overrides this at Inference.jl for transition::DynamicPPL.ParamsWithStats.

This override extracts transition.params and transition.stats directly—it doesn't call getparams(state) or getstats(state).

[shravanngoswamii]

Right now, the ParamsWithStats override:

• Extracts params from transition.params
• Extracts stats from transition.stats (includes logprior, loglikelihood, logjoint, and AHMC diagnostics like tree_depth, n_steps, acceptance_rate)
• Returns empty extras (always NamedTuple())
• Ignores state entirely

I am not sure if ignoring state is acceptable? Is the current design acceptable, or should we populate extras with static sampler configuration from state?
What are your thoughts @penelopeysm @yebai ?

[penelopeysm]

Since it's nonbreaking to add extra info in, I would be OK with this current implementation right now, and afterwards if there's a need for more information that can be changed later / in a separate PR.

[yebai]

I’m happy with the PR!

[penelopeysm]

Why the minor version bump?

[shravanngoswamii]

Why the minor version bump?

New funtion overloaded should be patch or minor?

[penelopeysm]

It's backwards compatible new functionality. If the version number is 0.x, it only needs to be a patch bump. If the version number is N.x (N ≥ 1) it only needs to be a minor bump.

But semver says you do have a choice: if you prefer to release it as a minor version 0.(x+1) or a major version (N+1), you can.

If you do choose to release as a minor version, please can you rebase it against breaking? See description in https://turinglang.org/docs/developers/contributing/#pull-requests-versions-and-releases (although we only really follow this for Turing and DPPL right now)

[penelopeysm]

I would recommend not doing it as a minor version though, as the next minor version will take a while to be done, since I have to do #2756 and that is very difficultl

[shravanngoswamii]

But semver says you do have a choice: if you prefer to release it as a minor version 0.(x+1) or a major version (N+1), you can.

Okay, I get it! Thanks for all the reviews, @penelopeysm. Let me know if you’re okay with merging this.