Fix indexing for chains in different threads

Author: penelopeysm

Project.toml

@@ -3,7 +3,7 @@ uuid = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 keywords = ["markov chain monte carlo", "probabilistic programming"]
 license = "MIT"
 desc = "A lightweight interface for common MCMC methods."
-version = "5.6.0"
+version = "5.6.1"
 
 [deps]
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"

src/sample.jl

@@ -391,7 +391,6 @@ function mcmcsample(
 
     # Copy the random number generator, model, and sample for each thread
     nchunks = min(nchains, Threads.nthreads())
-    chunksize = cld(nchains, nchunks)
     interval = 1:nchunks
     # `copy` instead of `deepcopy` for RNGs: https://github.com/JuliaLang/julia/issues/42899
     rngs = [copy(rng) for _ in interval]
@@ -434,15 +433,24 @@ function mcmcsample(
                 end
             end
 
+            # If nchains/nchunks = m with remainder n, then the first n chunks
+            # will have m + 1 chains, and the rest will have m chains
+            m, n = divrem(nchains, nchunks)
+            chain_index_groups = UnitRange{Int}[]
+            current_index = 1
+            for i in interval
+                nchains_this_chunk = i <= n ? m + 1 : m
+                push!(
+                    chain_index_groups,
+                    current_index:(current_index + nchains_this_chunk - 1),
+                )
+                current_index += nchains_this_chunk
+            end
+
             Distributed.@async begin
                 try
-                    Distributed.@sync for (i, _rng, _model, _sampler) in
-                                          zip(1:nchunks, rngs, models, samplers)
-                        chainidxs = if i == nchunks
-                            ((i - 1) * chunksize + 1):nchains
-                        else
-                            ((i - 1) * chunksize + 1):(i * chunksize)
-                        end
+                    Distributed.@sync for (chainidxs, _rng, _model, _sampler) in
+                                          zip(chain_index_groups, rngs, models, samplers)
                         Threads.@spawn for chainidx in chainidxs
                             # Seed the chunk-specific random number generator with the pre-made seed.
                             Random.seed!(_rng, seeds[chainidx])