Skip to content

mcmc() doesn't need to do trace() during warmup #834

Description

@njtierney

We do self$trace() during MCMC warmup, but we don't actually need those traced values, since we use other summaries.

The issue is that greta calls self$trace() once per burst in warmup, and appends that burst's free state onto traced_free_state with rbind. This ends up being a vector-growing memory issue - rbind copies everything accumulated so far on every append, so the cost grows with the square of the burst count.

Importantly, we don't even end up using the values we save in trace() in the warmup.

We've got some benchmarks of removing trace() in warmup at: greta.benchmarks/2026-08-20-warmup-trace.

git blame/explain

The trace() used to be there because it fed tuning, but then then this was replaced with using the welford accumulator for tuning:

  1. 92c876af (21 Mar 2018) added the warmup trace() call and the scrub, so tune_diag_sd() could do samples <- self$traced_free_state and take the sample posterior variance from it. Before this, warmup did not trace at all.
  2. 3c433f96 (4 Sep 2018) made these fields per-chain lists, so the append became the mapply(rbind, ...) still in trace() today and the read became a loop over chains.
  3. ef013050 (11 Sep 2018, "use welford accumulator for tuning") deleted that loop from tune_diag_sd() and replaced sample_variance(samples) with self$sample_variance(), backed by an online Welford accumulator fed from last_burst_free_states.

Note this is specifically tune_diag_sd() - trace_values() was untouched by that commit and still reads traced_free_state, which is the sampling-phase consumer and why the sampling trace() call must stay.

So it is a leftover from an R-side tuning refactor, not from TF1.

Fix

Remove the self$trace() call from the warmup loop in R/sampler_class.R. The sampling loop keeps its own.

Checked:

  • draws are bit-identical with and without it. Fixing both seeds (set.seed() and tf$random$set_seed() — see below) makes mcmc() reproducible, and 2 chains x 150 draws x 5 parameters come back identical, max absolute difference 0
  • no TF or TFP code can see it. With values = FALSE, trace() does one thing: mapply(rbind, ...) into an R6 field. Its input is already as.array()-converted R matrices. No tensor is touched, no Python round trip is made. The values = TRUE branch is the only one calling into TF, and warmup never takes it
  • the warmup progress bar is unaffected — it is driven by the burst loop's completed_iterations, and this removes the trace, not the bursts
  • aborting mid-warmup still returns NULL from stashed_samples(), which gates on traced_values; warmup calls trace() with values = FALSE, so it never fills that
  • full test suite unchanged, FAIL 0 | WARN 0 | SKIP 2 | PASS 1952

Relation to #547 and #765

Aside, relevant to #285 / #427

mcmc() is reproducible today if both seeds are set:

set.seed(1)
tf$random$set_seed(1L)

Neither alone is enough; together they give bit-identical draws, and different seeds give different draws. greta's own set_tf_seed() stores self$seed in dag$tf_environment$rng_seed, which nothing reads — it never calls tf$random$set_seed(). Relevant to #285 and #427

(Found during #745)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions