Merge branch 'breaking' into py/actually-use-init

penelopeysm · web-flow · commit 1bdb76ed55cd · 2025-09-15T16:56:45.000+01:00
diff --git a/.github/workflows/Benchmarking.yml b/.github/workflows/Benchmarking.yml
@@ -9,7 +9,7 @@ jobs:
 
     steps:
       - name: Checkout Repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           ref: ${{ github.event.pull_request.head.sha }}
 
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -56,7 +56,7 @@ jobs:
           - Group2
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
 
       - uses: julia-actions/setup-julia@v2
         with:
diff --git a/.github/workflows/DocTest.yml b/.github/workflows/DocTest.yml
@@ -26,7 +26,7 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
 
       - uses: julia-actions/setup-julia@v2
         with:
diff --git a/.github/workflows/DocsNav.yml b/.github/workflows/DocsNav.yml
@@ -17,7 +17,7 @@ jobs:
 
     steps:
       - name: Checkout gh-pages branch
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           ref: gh-pages
 
diff --git a/.github/workflows/DocsPreviewCleanup.yml b/.github/workflows/DocsPreviewCleanup.yml
@@ -9,7 +9,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout gh-pages branch
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           ref: gh-pages
       - name: Delete preview and history + push changes
diff --git a/.github/workflows/Enzyme.yml b/.github/workflows/Enzyme.yml
@@ -20,7 +20,7 @@ jobs:
   enzyme:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
 
       - uses: julia-actions/setup-julia@v2
         with:
diff --git a/.github/workflows/IntegrationTest.yml b/.github/workflows/IntegrationTest.yml
@@ -21,13 +21,13 @@ jobs:
           - {user: TuringLang, repo: Turing.jl}
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
           version: 1
       - uses: julia-actions/julia-buildpkg@v1
       - name: Clone Downstream
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           repository: ${{ matrix.package.user }}/${{ matrix.package.repo }}
           path: downstream
diff --git a/.github/workflows/JuliaPre.yml b/.github/workflows/JuliaPre.yml
@@ -23,7 +23,7 @@ jobs:
           - Group2
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
           version: 'pre' # pre-release
diff --git a/HISTORY.md b/HISTORY.md
@@ -2,7 +2,14 @@
 
 ## 0.38.0
 
-[...]
+The `varname_leaves` and `varname_and_value_leaves` functions have been moved to AbstractPPL.jl.
+Their behaviour is otherwise identical.
+
+## 0.37.2
+
+Make the `resume_from` keyword work for multiple-chain (parallel) sampling as well.
+Prior to this version, it was silently ignored.
+Note that to get the correct behaviour you also need to have a recent version of MCMCChains (v7.2.1).
 
 ## 0.37.1
 
diff --git a/Project.toml b/Project.toml
@@ -47,7 +47,7 @@ DynamicPPLMooncakeExt = ["Mooncake"]
 [compat]
 ADTypes = "1"
 AbstractMCMC = "5"
-AbstractPPL = "0.13"
+AbstractPPL = "0.13.1"
 Accessors = "0.1"
 BangBang = "0.4.1"
 Bijectors = "0.13.18, 0.14, 0.15"
diff --git a/docs/src/api.md b/docs/src/api.md
@@ -435,8 +435,6 @@ DynamicPPL.maybe_invlink_before_eval!!
 Base.merge(::AbstractVarInfo)
 DynamicPPL.subset
 DynamicPPL.unflatten
-DynamicPPL.varname_leaves
-DynamicPPL.varname_and_value_leaves
 ```
 
 ### Evaluation Contexts
diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl
@@ -1,12 +1,7 @@
 module DynamicPPLMCMCChainsExt
 
-if isdefined(Base, :get_extension)
-    using DynamicPPL: DynamicPPL
-    using MCMCChains: MCMCChains
-else
-    using ..DynamicPPL: DynamicPPL
-    using ..MCMCChains: MCMCChains
-end
+using DynamicPPL: DynamicPPL, AbstractPPL
+using MCMCChains: MCMCChains
 
 # Load state from a `Chains`: By convention, it is stored in `:samplerstate` metadata
 function DynamicPPL.loadstate(chain::MCMCChains.Chains)
@@ -136,7 +131,7 @@ function DynamicPPL.predict(
         varname_vals = mapreduce(
             collect,
             vcat,
-            map(DynamicPPL.varname_and_value_leaves, keys(vals), values(vals)),
+            map(AbstractPPL.varname_and_value_leaves, keys(vals), values(vals)),
         )
 
         return (varname_and_values=varname_vals, logp=DynamicPPL.getlogjoint(varinfo))
diff --git a/src/default_accumulators.jl b/src/default_accumulators.jl
@@ -165,9 +165,7 @@ function accumulate_observe!!(acc::LogLikelihoodAccumulator, right, left, vn)
     return acclogp(acc, Distributions.loglikelihood(right, left))
 end
 
-function default_accumulators(
-    ::Type{FloatT}=LogProbType, ::Type{IntT}=Int
-) where {FloatT,IntT}
+function default_accumulators(::Type{FloatT}=LogProbType) where {FloatT}
     return AccumulatorTuple(
         LogPriorAccumulator{FloatT}(),
         LogJacobianAccumulator{FloatT}(),
diff --git a/src/model_utils.jl b/src/model_utils.jl
@@ -81,7 +81,7 @@ function varname_in_chain!(
     # We use `VarName{sym}()` so that the resulting leaf `vn` only contains the tail of the optic.
     # This way we can use `getoptic(vn)` to extract the value from `x` and use `getoptic(vn) ∘ vn_parent`
     # to extract the value from the `chain`.
-    for vn in varname_leaves(VarName{sym}(), x)
+    for vn in AbstractPPL.varname_leaves(VarName{sym}(), x)
         # Update `out`, possibly in place, and return.
         l = AbstractPPL.getoptic(vn)
         varname_in_chain!(x, l ∘ vn_parent, chain, chain_idx, iteration_idx, out)
@@ -107,7 +107,7 @@ function values_from_chain(
     # This way we can use `getoptic(vn)` to extract the value from `x` and use `getoptic(vn) ∘ vn_parent`
     # to extract the value from the `chain`.
     out = similar(x)
-    for vn in varname_leaves(VarName{sym}(), x)
+    for vn in AbstractPPL.varname_leaves(VarName{sym}(), x)
         # Update `out`, possibly in place, and return.
         l = AbstractPPL.getoptic(vn)
         out = Accessors.set(
diff --git a/src/sampler.jl b/src/sampler.jl
@@ -84,6 +84,14 @@ function default_varinfo(::Random.AbstractRNG, ::Model, ::AbstractSampler)
     return typed_varinfo(VarInfo())
 end
 
+"""
+    init_strategy(sampler)
+
+Define the initialisation strategy used for generating initial values when
+sampling with `sampler`. Defaults to `InitFromPrior()`, but can be overridden.
+"""
+init_strategy(::Sampler) = InitFromPrior()
+
 function AbstractMCMC.sample(
     rng::Random.AbstractRNG,
     model::Model,
@@ -99,13 +107,22 @@ function AbstractMCMC.sample(
     )
 end
 
-"""
-    init_strategy(sampler)
-
-Define the initialisation strategy used for generating initial values when
-sampling with `sampler`. Defaults to `InitFromPrior()`, but can be overridden.
-"""
-init_strategy(::Sampler) = InitFromPrior()
+function AbstractMCMC.sample(
+    rng::Random.AbstractRNG,
+    model::Model,
+    sampler::Sampler,
+    parallel::AbstractMCMC.AbstractMCMCEnsemble,
+    N::Integer,
+    nchains::Integer;
+    chain_type=default_chain_type(sampler),
+    resume_from=nothing,
+    initial_state=loadstate(resume_from),
+    kwargs...,
+)
+    return AbstractMCMC.mcmcsample(
+        rng, model, sampler, parallel, N, nchains; chain_type, initial_state, kwargs...
+    )
+end
 
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
diff --git a/src/test_utils.jl b/src/test_utils.jl
@@ -11,7 +11,7 @@ using Bijectors: Bijectors
 using Accessors: Accessors
 
 # For backwards compat.
-using DynamicPPL: varname_leaves, update_values!!
+using DynamicPPL: update_values!!
 
 include("test_utils/model_interface.jl")
 include("test_utils/models.jl")
diff --git a/src/test_utils/ad.jl b/src/test_utils/ad.jl
@@ -180,12 +180,16 @@ Everything else is optional, and can be categorised into several groups:
    parameters using the `params` keyword argument.
 
 1. _Which type of logp is being calculated._
+
    By default, `run_ad` evaluates the 'internal log joint density' of the model,
    i.e., the log joint density in the unconstrained space. Thus, for example, in
+
        @model f() = x ~ LogNormal()
+
    the internal log joint density is `logpdf(Normal(), log(x))`. This is the
    relevant log density for e.g. Hamiltonian Monte Carlo samplers and is therefore
    the most useful to test.
+
    If you want the log joint density in the original model parameterisation, you
    can use `getlogjoint`. Likewise, if you want only the prior or likelihood,
    you can use `getlogprior` or `getloglikelihood`, respectively.
diff --git a/src/test_utils/sampler.jl b/src/test_utils/sampler.jl
@@ -51,7 +51,7 @@ function test_sampler(
         for vn in filter(varnames_filter, varnames(model))
             # We want to compare elementwise which can be achieved by
             # extracting the leaves of the `VarName` and the corresponding value.
-            for vn_leaf in varname_leaves(vn, get(target_values, vn))
+            for vn_leaf in AbstractPPL.varname_leaves(vn, get(target_values, vn))
                 target_value = get(target_values, vn_leaf)
                 chain_mean_value = marginal_mean_of_samples(chain, vn_leaf)
                 @test chain_mean_value ≈ target_value atol = atol rtol = rtol
diff --git a/src/utils.jl b/src/utils.jl
diff --git a/test/Project.toml b/test/Project.toml
diff --git a/test/contexts.jl b/test/contexts.jl
diff --git a/test/model.jl b/test/model.jl
diff --git a/test/model_utils.jl b/test/model_utils.jl
diff --git a/test/sampler.jl b/test/sampler.jl
diff --git a/test/test_util.jl b/test/test_util.jl
diff --git a/test/varinfo.jl b/test/varinfo.jl
