Fix non-reproducible step sizes (#1924)

* Fix non-reproducible step sizes * Add test
TuringLang · Dec 24, 2022 · f9353f0 · f9353f0 · yebai · Dec 24, 2022
1 parent 61b06f6
commit f9353f0
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.23.1"
+version = "0.23.2"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

diff --git a/src/inference/hmc.jl b/src/inference/hmc.jl
@@ -186,7 +186,7 @@ function DynamicPPL.initialstep(
 
     # Find good eps if not provided one
     if iszero(spl.alg.ϵ)
-        ϵ = AHMC.find_good_stepsize(hamiltonian, theta)
+        ϵ = AHMC.find_good_stepsize(rng, hamiltonian, theta)
         @info "Found initial step size" ϵ
     else
         ϵ = spl.alg.ϵ
@@ -550,7 +550,7 @@ function HMCState(
 
     # Find good eps if not provided one
     if iszero(spl.alg.ϵ)
-        ϵ = AHMC.find_good_stepsize(h, θ_init)
+        ϵ = AHMC.find_good_stepsize(rng, h, θ_init)
         @info "Found initial step size" ϵ
     else
         ϵ = spl.alg.ϵ

diff --git a/test/inference/hmc.jl b/test/inference/hmc.jl
@@ -207,4 +207,13 @@
         end
         @test sample(rng, mwe3(), HMC(0.2, 4), 1_000) isa Chains
     end
+
+    # issue #1923
+    @turing_testset "reproducibility" begin
+        alg = NUTS(1000, 0.8)
+        res1 = sample(StableRNG(123), gdemo_default, alg, 1000)
+        res2 = sample(StableRNG(123), gdemo_default, alg, 1000)
+        res3 = sample(StableRNG(123), gdemo_default, alg, 1000)
+        @test Array(res1) == Array(res2) == Array(res3)
+    end
 end