From 6d70f71d9b16f3a74f4ebdeda810daee4d480988 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Tue, 7 Oct 2025 13:04:34 +0100 Subject: [PATCH 1/2] Pass `num_warmup` kwarg to `step_warmup()` --- Project.toml | 2 +- src/sample.jl | 20 ++++++++++---------- test/utils.jl | 3 +++ 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/Project.toml b/Project.toml index f49b214a..cc7d8b8d 100644 --- a/Project.toml +++ b/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.8.1" +version = "5.8.2" [deps] BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" diff --git a/src/sample.jl b/src/sample.jl index 13e91b7a..ec772c0a 100644 --- a/src/sample.jl +++ b/src/sample.jl @@ -179,9 +179,9 @@ function mcmcsample( # Obtain the initial sample and state. sample, state = if num_warmup > 0 if initial_state === nothing - step_warmup(rng, model, sampler; kwargs...) + step_warmup(rng, model, sampler; num_warmup, kwargs...) else - step_warmup(rng, model, sampler, initial_state; kwargs...) + step_warmup(rng, model, sampler, initial_state; num_warmup, kwargs...) end else if initial_state === nothing @@ -202,7 +202,7 @@ function mcmcsample( for j in 1:discard_initial # Obtain the next sample and state. sample, state = if j ≤ num_warmup - step_warmup(rng, model, sampler, state; kwargs...) + step_warmup(rng, model, sampler, state; num_warmup, kwargs...) else step(rng, model, sampler, state; kwargs...) end @@ -229,7 +229,7 @@ function mcmcsample( for _ in 1:(thinning - 1) # Obtain the next sample and state. sample, state = if i ≤ keep_from_warmup - step_warmup(rng, model, sampler, state; kwargs...) + step_warmup(rng, model, sampler, state; num_warmup, kwargs...) else step(rng, model, sampler, state; kwargs...) end @@ -244,7 +244,7 @@ function mcmcsample( # Obtain the next sample and state. sample, state = if i ≤ keep_from_warmup - step_warmup(rng, model, sampler, state; kwargs...) + step_warmup(rng, model, sampler, state; num_warmup, kwargs...) else step(rng, model, sampler, state; kwargs...) end @@ -328,9 +328,9 @@ function mcmcsample( # Obtain the initial sample and state. sample, state = if num_warmup > 0 if initial_state === nothing - step_warmup(rng, model, sampler; kwargs...) + step_warmup(rng, model, sampler; num_warmup, kwargs...) else - step_warmup(rng, model, sampler, initial_state; kwargs...) + step_warmup(rng, model, sampler, initial_state; num_warmup, kwargs...) end else if initial_state === nothing @@ -344,7 +344,7 @@ function mcmcsample( for j in 1:discard_initial # Obtain the next sample and state. sample, state = if j ≤ num_warmup - step_warmup(rng, model, sampler, state; kwargs...) + step_warmup(rng, model, sampler, state; num_warmup, kwargs...) else step(rng, model, sampler, state; kwargs...) end @@ -364,7 +364,7 @@ function mcmcsample( for _ in 1:(thinning - 1) # Obtain the next sample and state. sample, state = if i ≤ keep_from_warmup - step_warmup(rng, model, sampler, state; kwargs...) + step_warmup(rng, model, sampler, state; num_warmup, kwargs...) else step(rng, model, sampler, state; kwargs...) end @@ -372,7 +372,7 @@ function mcmcsample( # Obtain the next sample and state. sample, state = if i ≤ keep_from_warmup - step_warmup(rng, model, sampler, state; kwargs...) + step_warmup(rng, model, sampler, state; num_warmup, kwargs...) else step(rng, model, sampler, state; kwargs...) end diff --git a/test/utils.jl b/test/utils.jl index b041b3a7..9ed92ba7 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -26,8 +26,11 @@ function AbstractMCMC.step_warmup( state::Union{Nothing,Integer}=nothing; loggers=false, initial_params=nothing, + num_warmup, kwargs..., ) + num_warmup isa Integer || + error("num_warmup should have been passed as a keyword argument to step_warmup") transition, state = AbstractMCMC.step( rng, model, sampler, state; loggers, initial_params, kwargs... ) From fdab17c889016432346c0765f3e3d70f779175a8 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Tue, 7 Oct 2025 13:10:07 +0100 Subject: [PATCH 2/2] Mention the num_warmup kwarg in step_warmup's docstring --- src/interface.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/interface.jl b/src/interface.jl index b58ced99..902424d2 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -83,6 +83,10 @@ When sampling using [`sample`](@ref), this takes the place of [`AbstractMCMC.ste This is useful if the sampler has an initial "warmup"-stage that is different from the standard iteration. +The total number of warmup steps requested in sampling will be passed to the `step_warmup` +function as the `num_warmup` keyword argument. This allows implementations of `step_warmup` +to customise their behavior based on this information. + By default, this simply calls [`AbstractMCMC.step`](@ref). """ step_warmup(rng, model, sampler; kwargs...) = step(rng, model, sampler; kwargs...)