diff --git a/src/sampler.jl b/src/sampler.jl index 1b282383..49e68073 100644 --- a/src/sampler.jl +++ b/src/sampler.jl @@ -228,12 +228,13 @@ function sample( end # Report end of sampling if verbose - EBFMI_est = EBFMI(map(s -> s.hamiltonian_energy, stats)) average_acceptance_rate = mean(map(s -> s.acceptance_rate, stats)) if θ isa AbstractVector n_chains = 1 + EBFMI_est = length(stats) > 1 ? EBFMI(map(s -> s.hamiltonian_energy, stats)) : NaN else n_chains = size(θ, 2) + EBFMI_est = length(stats) > 1 ? EBFMI(map(s -> s.hamiltonian_energy, stats)) : fill(NaN, n_chains) # Make sure that arrays are on CPU before printing. EBFMI_est = convert(Vector{eltype(EBFMI_est)}, EBFMI_est) average_acceptance_rate = convert( diff --git a/test/sampler.jl b/test/sampler.jl index 0c9e4e3d..9bfa2c3a 100644 --- a/test/sampler.jl +++ b/test/sampler.jl @@ -180,3 +180,14 @@ end @test length(stats) == n_samples end end + +@testset "n_samples=1 does not crash" begin + # Regression test for https://github.com/TuringLang/AdvancedHMC.jl/issues/281 + # verbose=true triggers EBFMI computation which failed on a 1-element stats vector + θ_init = rand(MersenneTwister(1), D) + h = Hamiltonian(DiagEuclideanMetric(D), ℓπ, ∂ℓπ∂θ) + κ = HMCKernel(Trajectory{MultinomialTS}(Leapfrog(0.1), GeneralisedNoUTurn())) + samples, stats = sample(h, κ, θ_init, 1; verbose=true, progress=false) + @test length(samples) == 1 + @test length(stats) == 1 +end