Insane memory-usage when using Emcee with large number of iterations (#1976)

* use `reduce` instead of splatting to avoid stack overflow

* added test for mem usage of Emcee

* version bump

* for emcee, walkers are now indexed along the chain index rather than
iteration index

* reduced number of iterations for emcee to see if that fixes tests

* attempt at reducing the memory usage of Emcee tests

* attempt at further reducing the memory footprint of the test

* fixed size issues

* Update emcee.jl

---------

Co-authored-by: Hong Ge <[email protected]>

Author: torfjelde

src/inference/emcee.jl

@@ -123,8 +123,13 @@ function AbstractMCMC.bundle_samples(
 
     # Extract names & construct param array.
     nms = [nms; extra_params]
-    parray = map(x -> hcat(x[1], x[2]), zip(vals_vec, extra_values_vec))
-    parray = cat(parray..., dims=3)
+    # `hcat` first to ensure we get the right `eltype`.
+    x = hcat(first(vals_vec), first(extra_values_vec))
+    # Pre-allocate to minimize memory usage.
+    parray = Array{eltype(x),3}(undef, length(vals_vec), size(x, 2), size(x, 1))
+    for (i, (vals, extras)) in enumerate(zip(vals_vec, extra_values_vec))
+        parray[i, :, :] = transpose(hcat(vals, extras))
+    end
 
     # Get the average or final log evidence, if it exists.
     le = getlogevidence(samples, state, spl)

test/inference/emcee.jl

@@ -1,5 +1,3 @@
-
-
 @testset "emcee.jl" begin
     @testset "gdemo" begin
         Random.seed!(9876)
@@ -14,6 +12,15 @@
         check_gdemo(chain)
     end
 
+    @testset "memory usage with large number of iterations" begin
+        # https://github.com/TuringLang/Turing.jl/pull/1976
+        @info "Testing emcee with large number of iterations"
+        spl = Emcee(10, 2.0)
+        n_samples = 10_000
+        chain = sample(gdemo_default, spl, n_samples)
+        check_gdemo(chain)
+    end
+
     @testset "initial parameters" begin
         nwalkers = 250
         spl = Emcee(nwalkers, 2.0)
@@ -30,7 +37,7 @@
 
         # Initial parameters
         chain = sample(gdemo_default, spl, 1; init_params=fill([2.0, 1.0], nwalkers))
-        @test chain[:s] == fill(2.0, nwalkers, 1)
-        @test chain[:m] == fill(1.0, nwalkers, 1)
+        @test chain[:s] == fill(2.0, 1, nwalkers)
+        @test chain[:m] == fill(1.0, 1, nwalkers)
     end
 end