`mcmc()` doesn't need to do trace() during warmup
- Dominant language
- C++
- Stars
- 607
- Forks
- 67
- Avg merge
- 3d 8h
- Merged PRs (30d)
- 1
Description
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](https://github.com/greta-dev/greta.benchmarks/tree/main/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`](https://github.com/greta-dev/greta/commit/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`](https://github.com/greta-dev/greta/commit/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`](https://github.com/greta-dev/greta/commit/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
- This is a piece of #547. If the warmup loop moves into TF, the R-side trace goes with it. Worth doing separately because it is a one-line deletion available now.
- #779, which resolves #765, does not fix it: it adds `do_warmup <- self$warmup > 0` so warmup can be skipped wholesale, and stubs `tune_tf` / `tune_r`, but never touches the `trace()` call.
# Aside, relevant to #285 / #427
`mcmc()` *is* reproducible today if both seeds are set:
```r
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)
Contributor guide
Research direction
Open R/sampler_class.R and locate the warmup loop's self$trace() call; compare it with the separate sampling-loop call that must remain. Remove only the warmup trace, then run the full test suite and check the warmup benchmark at greta.benchmarks/2026-08-20-warmup-trace; done means tests remain passing and sampling draws stay identical.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- r
- Domain
- backend
- Issue type
- Refactor
- Difficulty
- 1/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 90/100