sample with LogDensityFunction: part 1 - `hmc.jl`, `sghmc.jl`, `DynamicHMCExt` #2588

penelopeysm · 2025-06-08T13:43:27Z

This PR moves in the general direction of #2555.

It will take a long time to get everything to work, so I am trying to do this incrementally.

Summary

The fundamental idea (see #2555) is that we want

sample(::Union{Model, LogDensityFunction}, ::Union{InferenceAlgorithm, Sampler{<:InferenceAlgorithm}}, N)

to always forward to

sample(::LogDensityFunction, ::Sampler{<:InferenceAlgorithm}, N)

(along the way, we construct the LDF if we need to, and also construct the sampler if we need to).

Then, the concrete AbstractMCMC interface functions (i.e., mcmcsample, step) do not ever see a Model, they only see an LDF.

The future of `DynamicPPL.initialstep`

Note that this allows us to sidestep (and eventually, completely remove) DynamicPPL.initialstep. The reason why that function exists is because AbstractMCMC.step would do two things: first, generate the VarInfo that would eventually go into the LDF, and secondly, call initialstep (which was sampler-specific). Since the VarInfo generation bit is now handled in the LDF construction, it means that instead of having an extra function, we can just go back to implementing the two basic AbstractMCMC.step methods, which is a nice bonus.

Changes in this PR

Changing this all at once is bound to be not only impossible to do but also impossible to review. Thus, I've decided to (try to) implement this in a few stages. This PR is probably the one that makes the most sweeping changes, and also establishes the interface required. It:

Establishes the desired method dispatch behaviour for sample (see src/mcmc/abstractmcmc.jl). Because we aren't ready to extend this to every sampler and inference algorithm yet, these methods dispatch only on LDFCompatibleAlgorithm or LDFCompatibleSampler, which are defined at the top of the file. The idea is that we'll add samplers as we go along, and one day we'll eventually be ready to remove this type and just use InferenceAlgorithm.
When automatically constructing the LDF, there are a few things that we need to know to construct it properly:
- Does the VarInfo need to be linked?
- Does the LDF need to be constructed with an adtype?
This PR therefore also introduces interface functions that all (LDF-compatible) samplers must conform to, namely requires_unconstrained_space(::AbstractSampler) and get_adtype(::AbstractSampler). Sensible defaults of true and nothing are given. Note that these functions were already floating around the Turing codebase, so all I've really done is to bring it together and actually write docstrings for them.
Finally, there is an update_sample_kwargs function which samplers can use as a hook to modify the keyword arguments sent to sample(). See comments below for more details.

Fortunately

This doesn't actually require any changes to DynamicPPL, which I found to be a huge relief!

It's likely that some of the code in this PR will eventually be moved to DynamicPPL, as they don't have any non-DynamicPPL dependencies. But that can be handled very easily at a later stage, once we're confident that this all works.

Unfortunately

Changing the interface one sampler at a time completely breaks Gibbs, because for Gibbs to work, it requires all of its component samplers to be updated. So we may have to live with the Gibbs tests being broken for a while, and rely on me promising that I'll fix it at some point in time. In this PR, I've disabled the Gibbs tests that live outside test/mcmc/gibbs.jl.

Because I don't know how long this will take me, I don't even want to merge this into breaking, as I don't want to have a new release held up by the fact that only half the changes have been done. I've created a new base branch, sample-ldf, to collect all the work on this. When we're happy with it, we can merge that into breaking.

github-actions · 2025-06-08T13:45:17Z

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

codecov · 2025-06-08T15:07:19Z

Codecov Report

Attention: Patch coverage is 93.15068% with 10 lines in your changes missing coverage. Please review.

Project coverage is 80.49%. Comparing base (e84aec1) to head (8a0fb57).

Files with missing lines	Patch %	Lines
src/mcmc/hmc.jl	90.24%	4 Missing ⚠️
src/mcmc/abstractmcmc.jl	94.33%	3 Missing ⚠️
src/mcmc/algorithm.jl	33.33%	2 Missing ⚠️
src/mcmc/Inference.jl	83.33%	1 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (e84aec1) and HEAD (8a0fb57). Click for more details.

HEAD has 7 uploads less than BASE

Flag BASE (e84aec1) HEAD (8a0fb57)

28 21

Additional details and impacted files

@@              Coverage Diff               @@
##           sample-ldf    #2588      +/-   ##
==============================================
- Coverage       85.50%   80.49%   -5.02%     
==============================================
  Files              22       22              
  Lines            1456     1507      +51     
==============================================
- Hits             1245     1213      -32     
- Misses            211      294      +83

☔ 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 · 2025-06-08T15:03:04Z

ext/TuringDynamicHMCExt.jl

Many of the changes in sghmc.jl are quite similar to the ones in this file, so I added some comments explaining.

penelopeysm · 2025-06-08T15:03:59Z

ext/TuringDynamicHMCExt.jl

-struct DynamicNUTSState{L,V<:DynamicPPL.AbstractVarInfo,C,M,S}
-    logdensity::L
+struct DynamicNUTSState{V<:DynamicPPL.AbstractVarInfo,C,M,S}


The sampler state, traditionally, has included the LogDensityFunction as a field so that it doesn't need to be re-constructed on each iteration from the model + varinfo. This is no longer necessary because the LDF is itself an argument to AbstractMCMC.step.

penelopeysm · 2025-06-08T15:04:30Z

ext/TuringDynamicHMCExt.jl

-    # Ensure that initial sample is in unconstrained space.
-    if !DynamicPPL.islinked(vi)
-        vi = DynamicPPL.link!!(vi, model)
-        vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
-    end
-
-    # Define log-density function.
-    ℓ = DynamicPPL.LogDensityFunction(
-        model,
-        vi,
-        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
-        adtype=spl.alg.adtype,
-    )


All of this stuff is now handled inside AbstractMCMC.sample(), so there's no longer a need to duplicate this code inside every initialstep method.

penelopeysm · 2025-06-08T15:05:22Z

src/mcmc/Inference.jl

 function AbstractMCMC.bundle_samples(
    ts::Vector{<:Union{AbstractTransition,AbstractVarInfo}},
-    model::AbstractModel,
+    model_or_ldf::Union{DynamicPPL.Model,DynamicPPL.LogDensityFunction},
    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior,RepeatSampler},
    state,
    chain_type::Type{MCMCChains.Chains};


This signature is not super ideal, but it minimises breakage right now. Eventually when everything is fixed we can change the Union to just LDF.

penelopeysm · 2025-06-08T15:12:51Z

src/mcmc/hmc.jl

-# Handle setting `nadapts` and `discard_initial`
-function AbstractMCMC.sample(
-    rng::AbstractRNG,
-    model::DynamicPPL.Model,
-    sampler::Sampler{<:AdaptiveHamiltonian},
-    N::Integer;
-    chain_type=DynamicPPL.default_chain_type(sampler),
-    resume_from=nothing,
-    initial_state=DynamicPPL.loadstate(resume_from),
-    progress=PROGRESS[],
-    nadapts=sampler.alg.n_adapts,
-    discard_adapt=true,
-    discard_initial=-1,
-    kwargs...,
-)


The purpose of this overload was purely to modify the kwargs to sample. I ditched it in favour of adding a new hook, update_sample_kwargs, which does the same thing without abusing multiple dispatch. I think that function does the same thing. It's of course quite hard to prove this, although separating it into a different function does allow us to write unit tests for it to make sure that it's doing the right thing (which are now in test/mcmc/hmc.jl), so that's another benefit.

(Overloading sample() for individual samplers like this is quite precarious because we can't recursively call AbstractMCMC.sample or we will end up with infinite recursion -- it has to call mcmcsample. So, there's no way to 'extend' this with extra behaviour by e.g. calling another method of sample before calling mcmcsample.)

penelopeysm · 2025-06-08T15:14:25Z

src/mcmc/hmc.jl

-for alg in (:HMC, :HMCDA, :NUTS)
-    @eval getmetricT(::$alg{<:Any,metricT}) where {metricT} = metricT
-end
+getmetricT(::HMC{<:Any,metricT}) where {metricT} = metricT
+getmetricT(::HMCDA{<:Any,metricT}) where {metricT} = metricT
+getmetricT(::NUTS{<:Any,metricT}) where {metricT} = metricT


Metaprogramming is cool and all, but this wasn't really necessary, imo.

penelopeysm · 2025-06-08T15:17:57Z

test/mcmc/hmc.jl

        @testset "$(alg)" for alg in algs
            # Construct a HMC state by taking a single step
+            vi = DynamicPPL.VarInfo(gdemo_default)
+            vi = DynamicPPL.link(vi, gdemo_default)
+            ldf = LogDensityFunction(gdemo_default, vi; adtype=Turing.DEFAULT_ADTYPE)
            spl = Sampler(alg)
-            hmc_state = DynamicPPL.initialstep(
-                Random.default_rng(), gdemo_default, spl, DynamicPPL.VarInfo(gdemo_default)
-            )[2]
+            _, hmc_state = AbstractMCMC.step(Random.default_rng(), ldf, spl)


Finally, I think this test reveals one drawback of the current proposal: it becomes more annoying to directly call the AbstractMCMC interface. Let's say we want to benchmark the first step of a given sampler (for example, we were doing this the other day on the Gibbs sampler). Previously, we'd do:

rng = Random.default_rng() model = ... spl = ... @be AbstractMCMC.step(rng, model, spl)

Now, we have to do:

rng = Random.default_rng() model = ... vi = link(VarInfo(model), model) ldf = LogDensityFunction(model, vi; adtype=AutoForwardDiff()) spl = ... @be AbstractMCMC.step(rng, ldf, spl)

I think this is a fairly small price to pay because the occasions where we reach directly for AbstractMCMC interface are quite few, and the code simplification is more important than this. But I thought this was probably something just worth noting.

This would be less problematic if we introduced more convenient constructors for LDF: TuringLang/DynamicPPL.jl#863 so it might be worth keeping that in mind.

penelopeysm · 2025-06-17T12:38:19Z

@sunxd3 I removed the Sampler wrapper for HMC as you suggested, and I think adding the method for AbstractMCMC.step will have to come at the end (and we'll put it in DynamicPPL) but I've now added a note to #2555

sunxd3 · 2025-06-18T04:11:59Z

src/mcmc/Inference.jl

    @model,
    Metadata,
    VarInfo,
+    LogDensityFunction,
+    SimpleVarInfo,
+    AbstractVarInfo,


Being pedantic and annoying, I think it'll be better if we reorder so that the *VarInfos stay together.

sunxd3 · 2025-06-18T04:16:09Z

src/mcmc/algorithm.jl

+"""
+    update_sample_kwargs(spl::AbstractSampler, N::Integer, kwargs)
+
+Some samplers carry additional information about the keyword arguments that
+should be passed to `AbstractMCMC.sample`. This function provides a hook for
+them to update the default keyword arguments. The default implementation is for
+no changes to be made to `kwargs`.
+"""


I think it's worth adding that this function returns a NamedTuple (if I understand correctly).

also, it's prob worth mentioning what is N?

sunxd3 · 2025-06-18T04:22:16Z

src/mcmc/algorithm.jl

+get_adtype(::AbstractSampler) = nothing
+
+"""
+    requires_unconstrained_space(sampler::AbstractSampler)


Maybe this is a unimportant point: for something like MH, both constrained and unconstrained spaces might be fine.

On a high level, I do think "requires_unconstrained_space"-ness is a sampler property, but it's not sufficient condition to control linking.

Don't know if what I wrote makes sense.

sunxd3 · 2025-06-18T04:30:29Z

src/mcmc/hmc.jl

-abstract type Hamiltonian <: InferenceAlgorithm end
+# AbstractSampler interface for Turing
+
+abstract type Hamiltonian <: AbstractMCMC.AbstractSampler end


I want to tag @yebai for visibility.

I can see the motivation for InferenceAlgorithm, but I think this makes things a lot cleaner.

sunxd3 · 2025-06-18T04:39:02Z

src/mcmc/hmc.jl

    initial_params=nothing,
    nadapts=0,
    kwargs...,
 )
-    # Transform the samples to unconstrained space and compute the joint log probability.
-    vi = DynamicPPL.link(vi_original, model)
+    ldf.adtype === nothing &&


this is a good sanity check, does it make sense to force ldf and spl has the same adtype?

sunxd3 · 2025-06-18T05:00:49Z

src/mcmc/abstractmcmc.jl

+# This file contains the basic methods for `AbstractMCMC.sample`.
+# The overall aim is that users can call
+#
+#    sample(::Model, ::InferenceAlgorithm, N)


I assume InferenceAlgorithm here is still needed until we update all the interface?

github-actions bot assigned penelopeysm Jun 8, 2025

penelopeysm changed the base branch from sample-ldf to main June 8, 2025 13:43

penelopeysm changed the base branch from main to sample-ldf June 8, 2025 13:43

penelopeysm changed the title ~~sample with LogDensityFunction: part 1 - HMC~~ sample with LogDensityFunction: part 1 - hmc.jl + sghmc.jl Jun 8, 2025

penelopeysm changed the title ~~sample with LogDensityFunction: part 1 - hmc.jl + sghmc.jl~~ sample with LogDensityFunction: part 1 - hmc.jl, sghmc.jl, DynamicHMCExt Jun 8, 2025

penelopeysm commented Jun 8, 2025

View reviewed changes

penelopeysm force-pushed the py/ldf-hmc branch 2 times, most recently from 2d730bc to 1dcd37c Compare June 11, 2025 19:54

penelopeysm added 7 commits June 12, 2025 08:38

Fix imports (#2589)

8be9094

Bump DynamicPPL -> 0.36.8

480fb69

Add interface functions for InferenceAlgorithm

6599cfa

Replace _check_model with ordinary DynamicPPL.check_model

ad48370

Update AbstractMCMC interface for Hamiltonian samplers

ae7abde

Add new interface for AbstractMCMC.sample with LDFCompatibleAlgorithm

2c9f970

Fix existing tests

6923a72

penelopeysm force-pushed the py/ldf-hmc branch from 1dcd37c to 6923a72 Compare June 12, 2025 13:16

penelopeysm added 3 commits June 15, 2025 18:12

Add unit tests for HMC InferenceAlgorithm interface

4606bdd

Improve docstring

17a1f00

Add SamplingContext

71a8cf2

penelopeysm requested a review from sunxd3 June 16, 2025 14:16

penelopeysm mentioned this pull request Jun 16, 2025

sample with LogDensityFunction: part 2 - ess.jl + mh.jl #2590

Draft

penelopeysm added 2 commits June 17, 2025 00:00

Fix reproducibility of sample(ldf, ...) with deepcopy

85b1997

make Hamiltonian directly an AbstractSampler

49f6988

penelopeysm force-pushed the py/ldf-hmc branch from 3ff1c95 to 49f6988 Compare June 16, 2025 23:00

penelopeysm mentioned this pull request Jun 16, 2025

rogue idea: make SamplingContext a leaf context + make a new InitialisationContext TuringLang/DynamicPPL.jl#955

Open

penelopeysm added 2 commits June 17, 2025 10:26

Unwrap the other Hamiltonians to make them AbstractSamplers

e4cb590

Fix deepcopying

e08f548

penelopeysm mentioned this pull request Jun 17, 2025

Rework sample() call stack to use LogDensityFunction #2555

Open

10 tasks

remove unneeded Sampler tests

8a0fb57

sunxd3 reviewed Jun 18, 2025

View reviewed changes

sample with LogDensityFunction: part 1 - hmc.jl, sghmc.jl, DynamicHMCExt #2588

Are you sure you want to change the base?

sample with LogDensityFunction: part 1 - hmc.jl, sghmc.jl, DynamicHMCExt #2588

Uh oh!

Conversation

penelopeysm commented Jun 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

The future of DynamicPPL.initialstep

Changes in this PR

Fortunately

Unfortunately

Uh oh!

github-actions bot commented Jun 8, 2025

Uh oh!

codecov bot commented Jun 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

penelopeysm Jun 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

penelopeysm Jun 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

penelopeysm commented Jun 17, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sunxd3 Jun 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sunxd3 Jun 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sample with LogDensityFunction: part 1 - `hmc.jl`, `sghmc.jl`, `DynamicHMCExt` #2588

sample with LogDensityFunction: part 1 - `hmc.jl`, `sghmc.jl`, `DynamicHMCExt` #2588

penelopeysm commented Jun 8, 2025 •

edited

Loading

The future of `DynamicPPL.initialstep`

codecov bot commented Jun 8, 2025 •

edited

Loading

penelopeysm Jun 8, 2025 •

edited

Loading

penelopeysm Jun 8, 2025 •

edited

Loading

sunxd3 Jun 18, 2025 •

edited

Loading

sunxd3 Jun 18, 2025 •

edited

Loading