Fix non-reproducible step sizes (#1924)

* Fix non-reproducible step sizes

* Add test

Author: devmotion

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"

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.ϵ

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