Remove resume_from

Author: penelopeysm

HISTORY.md

@@ -50,6 +50,11 @@ Other functions such as `tilde_assume` and `assume` (and their `observe` counter
 Note that this was effectively already the case in DynamicPPL 0.37 (where they were just wrappers around each other).
 The separation of these functions was primarily implemented to avoid performing extra work where unneeded (e.g. to not calculate the log-likelihood when `PriorContext` was being used). This functionality has since been replaced with accumulators (see the 0.37 changelog for more details).
 
+### Removal of `resume_from`
+
+The `resume_from=chn` keyword argument to `sample` has been removed; please use `initial_state=DynamicPPL.loadstate(chn)` instead.
+`loadstate` is exported from DynamicPPL.
+
 **Other changes**
 
 ### Reimplementation of functions using `InitContext`

src/DynamicPPL.jl

@@ -129,6 +129,8 @@ export AbstractVarInfo,
     prefix,
     returned,
     to_submodel,
+    # Chain save/resume
+    loadstate,
     # Convenience macros
     @addlogprob!,
     value_iterator_from_chain,

src/sampler.jl

@@ -58,14 +58,10 @@ function AbstractMCMC.sample(
     model::Model,
     sampler::Sampler,
     N::Integer;
-    chain_type=default_chain_type(sampler),
-    resume_from=nothing,
-    initial_state=loadstate(resume_from),
+    initial_state=nothing,
     kwargs...,
 )
-    return AbstractMCMC.mcmcsample(
-        rng, model, sampler, N; chain_type, initial_state, kwargs...
-    )
+    return AbstractMCMC.mcmcsample(rng, model, sampler, N; initial_state, kwargs...)
 end
 
 function AbstractMCMC.sample(
@@ -75,13 +71,11 @@ function AbstractMCMC.sample(
     parallel::AbstractMCMC.AbstractMCMCEnsemble,
     N::Integer,
     nchains::Integer;
-    chain_type=default_chain_type(sampler),
-    resume_from=nothing,
-    initial_state=loadstate(resume_from),
+    initial_state=nothing,
     kwargs...,
 )
     return AbstractMCMC.mcmcsample(
-        rng, model, sampler, parallel, N, nchains; chain_type, initial_state, kwargs...
+        rng, model, sampler, parallel, N, nchains; initial_state, kwargs...
     )
 end
 
@@ -107,20 +101,12 @@ function AbstractMCMC.step(
 end
 
 """
-    loadstate(data)
-
-Load sampler state from `data`.
-
-By default, `data` is returned.
-"""
-loadstate(data) = data
-
-"""
-    default_chain_type(sampler)
+    loadstate(chain::AbstractChains)
 
-Default type of the chain of posterior samples from `sampler`.
+Load sampler state from an `AbstractChains` object. This function should be overloaded by a
+concrete Chains implementation.
 """
-default_chain_type(::Sampler) = Any
+function loadstate end
 
 """
     initialstep(rng, model, sampler, varinfo; kwargs...)

test/sampler.jl

@@ -12,7 +12,7 @@
         @test AbstractMCMC.step(Xoshiro(468), g(), spl) isa Any
     end
 
-    @testset "initial_state and resume_from kwargs" begin
+    @testset "initial_state" begin
         # Model is unused, but has to be a DynamicPPL.Model otherwise we won't hit our
         # overloaded method.
         @model f() = x ~ Normal()
@@ -58,19 +58,10 @@
                 spl,
                 N_iters;
                 progress=false,
-                initial_state=chn.info.samplerstate,
+                initial_state=DynamicPPL.loadstate(chn),
                 chain_type=MCMCChains.Chains,
             )
             @test all(chn2[:x] .== initial_value)
-            # using `resume_from`
-            chn3 = sample(
-                model,
-                spl,
-                N_iters;
-                progress=false,
-                resume_from=chn,
-                chain_type=MCMCChains.Chains,
-            )
             @test all(chn3[:x] .== initial_value)
         end
 
@@ -94,21 +85,10 @@
                 N_iters,
                 N_chains;
                 progress=false,
-                initial_state=chn.info.samplerstate,
+                initial_state=DynamicPPL.loadstate(chn),
                 chain_type=MCMCChains.Chains,
             )
             @test all(i -> chn2[:x][i, :] == initial_value, 1:N_iters)
-            # using `resume_from`
-            chn3 = sample(
-                model,
-                spl,
-                MCMCThreads(),
-                N_iters,
-                N_chains;
-                progress=false,
-                resume_from=chn,
-                chain_type=MCMCChains.Chains,
-            )
             @test all(i -> chn3[:x][i, :] == initial_value, 1:N_iters)
         end
     end