Fix init_params=nothing (#401)

Fixes the test errors in TuringLang/Turing.jl#1799.

Co-authored-by: David Widmann <[email protected]>

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.19.0"
+version = "0.19.1"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/sampler.jl

@@ -69,7 +69,12 @@ end
 
 # initial step: general interface for resuming and
 function AbstractMCMC.step(
-    rng::Random.AbstractRNG, model::Model, spl::Sampler; resume_from=nothing, kwargs...
+    rng::Random.AbstractRNG,
+    model::Model,
+    spl::Sampler;
+    resume_from=nothing,
+    init_params=nothing,
+    kwargs...,
 )
     if resume_from !== nothing
         state = loadstate(resume_from)
@@ -81,8 +86,8 @@ function AbstractMCMC.step(
     vi = VarInfo(rng, model, _spl)
 
     # Update the parameters if provided.
-    if haskey(kwargs, :init_params)
-        vi = initialize_parameters!!(vi, kwargs[:init_params], spl)
+    if init_params !== nothing
+        vi = initialize_parameters!!(vi, init_params, spl)
 
         # Update joint log probability.
         # TODO: fix properly by using sampler and evaluation contexts
@@ -96,7 +101,7 @@ function AbstractMCMC.step(
         end
     end
 
-    return initialstep(rng, model, spl, vi; kwargs...)
+    return initialstep(rng, model, spl, vi; init_params=init_params, kwargs...)
 end
 
 """

test/sampler.jl

@@ -125,6 +125,26 @@
                 @test !ismissing(c[1].metadata.s.vals[1])
                 @test c[1].metadata.m.vals == [-1]
             end
+
+            # specify `init_params=nothing`
+            Random.seed!(1234)
+            chain1 = sample(model, sampler, 1; progress=false)
+            Random.seed!(1234)
+            chain2 = sample(model, sampler, 1; init_params=nothing, progress=false)
+            @test chain1[1].metadata.m.vals == chain2[1].metadata.m.vals
+            @test chain1[1].metadata.s.vals == chain2[1].metadata.s.vals
+
+            # parallel sampling
+            Random.seed!(1234)
+            chains1 = sample(model, sampler, MCMCThreads(), 1, 10; progress=false)
+            Random.seed!(1234)
+            chains2 = sample(
+                model, sampler, MCMCThreads(), 1, 10; init_params=nothing, progress=false
+            )
+            for (c1, c2) in zip(chains1, chains2)
+                @test c1[1].metadata.m.vals == c2[1].metadata.m.vals
+                @test c1[1].metadata.s.vals == c2[1].metadata.s.vals
+            end
         end
     end
 end