fixed initial_params and initial_state for MCMCDistributed

Author: torfjelde

Project.toml

@@ -9,6 +9,7 @@ version = "4.5.0"
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 ConsoleProgressMonitor = "88cd18e8-d9cc-4ea6-8889-5259c0d15c8b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
@@ -21,6 +22,7 @@ Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999"
 [compat]
 BangBang = "0.3.19"
 ConsoleProgressMonitor = "0.1"
+FillArrays = "1"
 LogDensityProblems = "2"
 LoggingExtras = "0.4, 0.5, 1"
 ProgressLogging = "0.1"

src/AbstractMCMC.jl

@@ -8,6 +8,7 @@ using ProgressLogging: ProgressLogging
 using StatsBase: StatsBase
 using TerminalLoggers: TerminalLoggers
 using Transducers: Transducers
+using FillArrays: FillArrays
 
 using Distributed: Distributed
 using Logging: Logging

src/sample.jl

@@ -432,6 +432,9 @@ function mcmcsample(
     check_initial_params(initial_params, nchains)
     check_initial_state(initial_state, nchains)
 
+    _initial_params = initial_params === nothing ? FillArrays.Fill(nothing, nchains) : initial_params
+    _initial_state = initial_state === nothing ? FillArrays.Fill(nothing, nchains) : initial_state
+
     # Create a seed for each chain using the provided random number generator.
     seeds = rand(rng, UInt, nchains)
 
@@ -490,7 +493,7 @@ function mcmcsample(
                         return chain
                     end
                     chains = Distributed.pmap(
-                        sample_chain, pool, seeds, initial_params, initial_state
+                        sample_chain, pool, seeds, _initial_params, _initial_state
                     )
                 finally
                     # Stop updating the progress bar.