Skip to content

Pass information through to step / step_warmup about whether a step will be discardedΒ #185

@penelopeysm

Description

@penelopeysm

There are two situations in which we call step / step_warmup, and then proceed to throw away the transition object and only use the state:

  1. initial discarded steps (burn-in)
  2. thinning

# Discard initial samples.
for j in 1:discard_initial
# Obtain the next sample and state.
sample, state = if j ≀ num_warmup
step_warmup(rng, model, sampler, state; num_warmup, kwargs...)
else
step(rng, model, sampler, state; kwargs...)
end

for _ in 1:(thinning - 1)
# Obtain the next sample and state.
sample, state = if i ≀ keep_from_warmup
step_warmup(rng, model, sampler, state; num_warmup, kwargs...)
else
step(rng, model, sampler, state; kwargs...)
end

We should pass an extra keyword argument into step and step_warmup to let implementations know whether they need to return a meaningful transition or not. This could have substantial performance benefits e.g. for Turing, where calculating the transition object requires model re-evaluation.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions