Remove resume_from argument

TuringLang · Jan 8, 2025 · a0e6f32 · a0e6f32
1 parent 003ff2f
commit a0e6f32
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 5 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.33.0"
+version = "0.34.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/docs/src/api.md b/docs/src/api.md
@@ -134,6 +134,12 @@ returned(::Model)
 
 ## Utilities
 
+To retrieve the final sampler state from a chain of samples (useful for resuming sampling from a previous point):
+
+```@docs
+loadstate
+```
+
 It is possible to manually increase (or decrease) the accumulated log density from within a model function.
 
 ```@docs
@@ -425,7 +431,6 @@ The default implementation of [`Sampler`](@ref) uses the following unexported fu
 
 ```@docs
 DynamicPPL.initialstep
-DynamicPPL.loadstate
 DynamicPPL.initialsampler
 ```
 

diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
@@ -92,6 +92,7 @@ export AbstractVarInfo,
     Sampler,
     SampleFromPrior,
     SampleFromUniform,
+    loadstate,
     # Contexts
     SamplingContext,
     DefaultContext,

diff --git a/src/sampler.jl b/src/sampler.jl
@@ -100,10 +100,16 @@ function AbstractMCMC.sample(
     sampler::Sampler,
     N::Integer;
     chain_type=default_chain_type(sampler),
-    resume_from=nothing,
-    initial_state=loadstate(resume_from),
+    initial_state=nothing,
     kwargs...,
 )
+    if haskey(kwargs, :resume_from)
+        throw(
+            ArgumentError(
+                "The `resume_from` keyword argument is no longer supported. Please use `initial_state=loadstate(chain)` instead of `resume_from=chain`.",
+            ),
+        )
+    end
     return AbstractMCMC.mcmcsample(
         rng, model, sampler, N; chain_type, initial_state, kwargs...
     )
@@ -135,7 +141,14 @@ end
 
 Load sampler state from `data`.
 
-By default, `data` is returned.
+If `data` isa MCMCChains.Chains object, this attempts to fetch the last state
+of the sampler from the metadata stored inside the Chains object. This requires
+you to have passed the `save_state=true` keyword argument to the `sample()`
+when generating the chain.
+
+This function can be overloaded for specific types of `data` if desired. If
+there is no specific implementation for a given type, it falls back to just
+returning `data`, i.e. acts as an identity function.
 """
 loadstate(data) = data