greta#853 (windowed-adaptation-i853): warmup schemes compared

Five warmup schemes run inside TensorFlow on five models: greta’s own, greta’s with greta#841’s count fixed, Stan’s windowed scheme from greta#853, and TFP’s windowed HMC and NUTS. Stan’s scheme finds the scales of the normals that greta’s own misses, and gives more effective samples per second than greta’s on every model but eight_schools.
Published

October 7, 2026

Setup
# the packages' start-up messages, such as where here() finds the project, are
# not part of the results
suppressPackageStartupMessages({
  library(dplyr)
  library(tidyr)
  library(ggplot2)
  library(bench)
  library(here)
})

post_path <- "posts/2026-10-07-warmup-adaptation-i853"

# each `# ---- label ----` section of the scripts, shown below by its label
knitr::read_chunk(here(post_path, "benchmark.R"))
knitr::read_chunk(here(post_path, "run.R"))

# links to the scripts on GitHub, at the commit being rendered when GitHub
# Actions renders the site
github_ref <- Sys.getenv("GITHUB_SHA", "main")
github_file_url <- function(file) {
  paste0(
    "https://github.com/greta-dev/greta.benchmarks/blob/",
    github_ref,
    "/",
    post_path,
    "/",
    file
  )
}

# a link to the lines of a script's `# ---- label ----` section. Quarto escapes
# the markdown in an inline `{r}` result unless it is wrapped in I()
section_link <- function(file, label) {
  lines <- readLines(here(post_path, file))
  headings <- grep("^# ---- .+ ----$", lines)
  start <- headings[lines[headings] == paste("# ----", label, "----")]
  end <- c(headings, length(lines) + 1)[match(start, headings) + 1] - 1
  while (end > start && !nzchar(trimws(lines[end]))) {
    end <- end - 1
  }
  I(paste0(
    "[`", label, "`](", github_file_url(file), "#L", start, "-L", end, ")"
  ))
}

# "a, b and c"
paste_and <- function(x) {
  if (length(x) < 2) {
    return(x)
  }
  paste(paste(head(x, -1), collapse = ", "), "and", tail(x, 1))
}

# links to several sections of a script, as "a, b and c"
section_links <- function(file, labels) {
  links <- vapply(labels, \(label) section_link(file, label), character(1))
  I(paste_and(links))
}

# three significant figures each, where format() would pad a vector to the
# decimal places of its smallest value. formatC() pads a number with fewer
# than three significant figures with spaces, which trimws() removes
format_number <- function(x) {
  trimws(formatC(signif(x, 3), format = "fg", digits = 3, big.mark = ","))
}

# "3.25 to 4.81": the smallest and largest of x
format_range <- function(x) {
  paste(format_number(min(x)), "to", format_number(max(x)))
}

# a searchable table, its numbers rounded for reading
dt_table <- function(data, digits = 3) {
  table <- DT::datatable(
    data,
    rownames = FALSE,
    filter = "top",
    options = list(pageLength = 10, autoWidth = TRUE)
  )
  decimal_columns <- names(data)[vapply(data, is.double, logical(1))]
  if (length(decimal_columns) == 0) {
    return(table)
  }
  DT::formatRound(table, decimal_columns, digits)
}

theme_set(theme_bw(base_size = 13))

sessions <- list.files(
  here(post_path, "results"),
  pattern = "^session-[0-9]+[.]rds$",
  full.names = TRUE
) |>
  lapply(readRDS)
n_sessions <- length(sessions)
results <- lapply(sessions, \(session_file) session_file$results$branch)
bench_settings <- results[[1]]$bench_settings
model_levels <- names(results[[1]]$schemes)
method_levels <- c(
  bench_settings$schemes,
  paste0("tfp_", bench_settings$tfp_kinds)
)
method_colours <- setNames(
  c("#999999", "#56B4E9", "#009E73", "#E69F00", "#CC79A7"),
  method_levels
)
method_scale <- scale_colour_manual(values = method_colours, guide = "none")
method_fill <- scale_fill_manual(values = method_colours, guide = "none")
jitter <- position_jitter(width = 0, height = 0.2, seed = 2026)

# each model's runs are its `bench-` section of benchmark.R
bench_links <- section_links(
  "benchmark.R",
  paste0("bench-", gsub("_", "-", model_levels))
)

# A raincloud for each method: the density of its sessions above, their box
# plot in the middle, and every session as a dot below. The summaries are
# drawn first and faint, so the dots sit on top.
raincloud <- list(
  ggdist::stat_halfeye(
    aes(fill = method),
    width = 0.4,
    justification = -0.3,
    .width = 0,
    point_colour = NA,
    alpha = 0.3
  ),
  geom_boxplot(
    aes(fill = method),
    width = 0.25,
    outlier.shape = NA,
    alpha = 0.3
  ),
  # a fixed dot size, so a session looks the same size in every panel
  ggdist::stat_dots(
    aes(fill = method),
    side = "bottom",
    layout = "swarm",
    justification = 1.3,
    scale = 0.75,
    binwidth = unit(2.6, "mm"),
    overflow = "compress",
    colour = NA,
    alpha = 0.6
  ),
  method_scale,
  method_fill,
  theme_bw(base_size = 13)
)

# one row per timed call, from every session, model and method
every_call <- lapply(seq_along(results), function(i) {
  lapply(model_levels, function(model_name) {
    schemes <- results[[i]]$schemes[[model_name]]$marks
    tfp <- results[[i]]$tfp[[model_name]]$marks
    marks <- list(
      tibble(method = schemes$scheme, time = schemes$time),
      tibble(method = paste0("tfp_", tfp$kind), time = tfp$time)
    )
    bind_rows(marks) |>
      mutate(call = lapply(time, seq_along)) |>
      unnest(c(time, call)) |>
      mutate(
        session = sessions[[i]]$session,
        model = model_name,
        seconds = as.numeric(time),
        .keep = "unused"
      )
  }) |>
    bind_rows()
}) |>
  bind_rows() |>
  mutate(
    method = factor(method, levels = method_levels),
    model = factor(model, levels = model_levels),
    calls = factor(
      ifelse(call == 1, "first call", "second call"),
      levels = c("first call", "second call")
    )
  )

# posterior caps a bulk ESS at S × log10(S) for S draws, and warns each time it
# does. Count those warnings here, keep the cap beside each estimate, and say
# what they mean under the effective samples plot.
ess_cap_warnings <- 0
summarise_counting_caps <- function(draws) {
  withCallingHandlers(
    posterior::summarise_draws(draws, "mean", "sd", "rhat", "ess_bulk"),
    warning = function(w) {
      if (grepl("The ESS has been capped", conditionMessage(w), fixed = TRUE)) {
        ess_cap_warnings <<- ess_cap_warnings + 1
        invokeRestart("muffleWarning")
      }
    }
  )
}

# a posterior summary of every method's second-call draws, in every session
posterior_summaries <- lapply(seq_along(results), function(i) {
  lapply(model_levels, function(model_name) {
    draws <- c(
      results[[i]]$schemes[[model_name]]$draws,
      setNames(
        results[[i]]$tfp[[model_name]]$draws,
        paste0("tfp_", names(results[[i]]$tfp[[model_name]]$draws))
      )
    )
    lapply(names(draws), function(method) {
      n_draws <- posterior::ndraws(draws[[method]])
      tibble(
        session = sessions[[i]]$session,
        model = model_name,
        method = method,
        summarise_counting_caps(draws[[method]]),
        n_draws = n_draws,
        ess_cap = n_draws * log10(n_draws)
      )
    }) |>
      bind_rows()
  }) |>
    bind_rows()
}) |>
  bind_rows() |>
  mutate(
    method = factor(method, levels = method_levels),
    model = factor(model, levels = model_levels),
    ess_capped = ess_bulk >= ess_cap * (1 - 1e-9)
  )

second_calls <- every_call |>
  filter(call == 2) |>
  select(session, model, method, seconds)

# each run's smallest bulk ESS and largest R-hat, and that ESS per second of
# the run's second call, warmup included
efficiency <- posterior_summaries |>
  group_by(session, model, method) |>
  summarise(
    smallest_ess_bulk = min(ess_bulk),
    smallest_ess_capped = ess_capped[which.min(ess_bulk)],
    largest_rhat = max(rhat),
    .groups = "drop"
  ) |>
  left_join(second_calls, by = c("session", "model", "method")) |>
  mutate(smallest_ess_per_second = smallest_ess_bulk / seconds)

# each method's median over sessions, as a multiple of greta's own scheme's
# median on the same model
relative_to_greta <- function(data, value) {
  data |>
    group_by(model, method) |>
    summarise(median = median({{ value }}), .groups = "drop") |>
    group_by(model) |>
    mutate(times_greta = median / median[method == "greta"]) |>
    ungroup()
}

# one method's multiples of greta's, named by model
times_greta <- function(ratios, scheme) {
  scheme_ratios <- filter(ratios, method == scheme)
  setNames(scheme_ratios$times_greta, scheme_ratios$model)
}

ess_per_second_ratios <- relative_to_greta(efficiency, smallest_ess_per_second)
windowed_ess_per_second <- times_greta(ess_per_second_ratios, "windowed")

greta#853 is an issue, not a PR: it proposes that hmc() adapt diag_sd in windows during warmup, as Stan does, and the windowed-adaptation-i853 branch adds that scheme to greta’s sampler as an internal choice. This post runs it beside greta’s own scheme, greta’s with greta#841’s count fixed, and TFP’s windowed HMC and NUTS, all inside TensorFlow on one greta commit and five models; it is not part of PR greta#855. The windowed scheme gives 275 times the smallest bulk ESS per second of greta’s own scheme on twenty normals of very different scales, 3.25 to 4.81 times on linear, multiple_linear and hierarchical_linear, and 0.729 times on eight_schools.

The five schemes

  • greta: greta’s own tuning. It adapts the step size, epsilon, by dual averaging over the first 10% and last 60% of warmup (Hoffman and Gelman 2014), and estimates diag_sd between 10% and 40% of warmup from every warmup draw so far, shrunk by a count of the rejected proposals, then fixes it.
  • greta_841: the same, shrunk by a count of the accepted proposals, as greta#841 says it was meant to be.
  • windowed: Stan’s windowed adaptation, from greta#853. After a buffer that adapts epsilon only, it estimates diag_sd afresh in windows ending at 25%, 45% and 90% of warmup, each from its own draws, and restarts dual averaging after each.
  • tfp_hmc and tfp_nuts: TFP’s own windowed adaptation of HMC and NUTS, as its windowed_adaptive_hmc() and windowed_adaptive_nuts() use it, on the same log density. TFP’s HMC takes 8 leapfrog steps every iteration, where greta’s draws a number between Lmin and Lmax.

Does each scheme find the normals’ scales?

For greta’s three schemes, each normal’s tuned diag_sd over its true standard deviation, so 1 is a perfect estimate. TFP keeps its estimate inside its kernel results, which the benchmark does not return.

Code
true_sd <- results[[1]]$normals_sd
tuned_scales <- lapply(seq_along(results), function(i) {
  tuned <- results[[i]]$schemes$normals$tuned
  lapply(names(tuned), function(scheme) {
    tibble(
      session = sessions[[i]]$session,
      method = scheme,
      true_sd = true_sd,
      ratio = tuned[[scheme]]$diag_sd / true_sd
    )
  }) |>
    bind_rows()
}) |>
  bind_rows() |>
  mutate(method = factor(method, levels = method_levels))

# each scheme's tuned diag_sd over the true sd, named by scheme
tuned_ratios <- split(tuned_scales$ratio, tuned_scales$method, drop = TRUE)
Code
ggplot(tuned_scales, aes(x = true_sd, y = ratio, colour = method)) +
  geom_hline(yintercept = 1, colour = "grey60") +
  geom_point(alpha = 0.7) +
  scale_x_log10() +
  scale_y_log10() +
  method_scale +
  facet_wrap(~method) +
  labs(x = "True standard deviation", y = "Tuned diag_sd / true sd")

Each dot is one normal in one session, on log scales. The grey line is a perfect estimate.

Tuned in benchmark.R’s bench-normals section, by run_schemes() from its benchmark-schemes section, which keeps each scheme’s diag_sd.

The windowed scheme’s diag_sd is 0.927 to 1.09 times the true standard deviation. greta’s own is 0.0115 to 1.47 times, furthest off on the widest normals, and greta#841’s count leaves it much the same, at 0.0117 to 1.4 times.

How long does each scheme take?

Seconds for the second call on each model, which reuses the first call’s trace. The effective samples per second below divide by these.

Code
speed_ratios <- relative_to_greta(second_calls, seconds)
Code
ggplot(second_calls, aes(x = seconds, y = method, colour = method)) +
  raincloud +
  scale_x_log10() +
  facet_wrap(~model, ncol = 1, scales = "free_x") +
  labs(x = "Seconds", y = NULL)

Each dot is one session’s second call, on a log scale. Above the dots, the curve is the density of the sessions and the box marks their median and quartiles.

Timed by benchmark.R’s bench-normals, bench-linear, bench-multiple-linear, bench-hierarchical-linear and bench-eight-schools sections, which call run_schemes() from its benchmark-schemes section and run_tfp() from its benchmark-tfp section.

On the five models, the windowed scheme takes 1.02 to 1.05 times as long as greta’s own, and greta_841 0.99 to 1.2 times, so the three cost about the same. TFP’s HMC takes 0.693 to 0.819 times as long as greta’s, and TFP’s NUTS 5.33 to 26.8 times.

Table: median seconds of the first and second calls
Code
every_call |>
  group_by(model, calls, method) |>
  summarise(seconds = median(seconds), .groups = "drop") |>
  pivot_wider(names_from = method, values_from = seconds) |>
  dt_table(digits = 2)

How many effective samples per second does each give?

The bulk ESS of each run’s worst-mixing parameter, as Hoffman and Gelman (2014) report it, per second of the second call, warmup included (Girolami and Calderhead 2011).

Code
capped_estimates <- filter(posterior_summaries, ess_capped)
capped_runs <- distinct(capped_estimates, session, model, method)

# every run as a small dot over a large faint mean, for the smallest ESS here
# and the largest R-hat in the checks
efficiency_means <- efficiency |>
  group_by(model, method) |>
  summarise(
    across(c(smallest_ess_bulk, largest_rhat), mean),
    .groups = "drop"
  )

gg_sessions <- function(x, x_label) {
  ggplot(efficiency, aes(x = {{ x }}, y = method, colour = method)) +
    geom_point(data = efficiency_means, size = 5, alpha = 0.25) +
    geom_point(position = jitter, size = 1.2) +
    method_scale +
    facet_wrap(~model, ncol = 1, scales = "free_x") +
    labs(x = x_label, y = NULL)
}
Code
ggplot(
  efficiency,
  aes(x = smallest_ess_per_second, y = method, colour = method)
) +
  raincloud +
  scale_x_log10() +
  facet_wrap(~model, ncol = 1, scales = "free_x") +
  labs(x = "Smallest bulk ESS per second", y = NULL)

Each dot is one session’s run, on a log scale. Above the dots, the curve is the density of the sessions and the box marks their median and quartiles.

Draws and times from benchmark.R’s bench-normals, bench-linear, bench-multiple-linear, bench-hierarchical-linear and bench-eight-schools sections, which call run_schemes() from its benchmark-schemes section and run_tfp() from its benchmark-tfp section.

The windowed scheme gives 275 times the smallest bulk ESS per second of greta’s own scheme on the normals, 3.25 to 4.81 times on linear, multiple_linear and hierarchical_linear, and 0.729 times on eight_schools. On the five models, greta_841 gives 0.83 to 1.45 times greta’s, TFP’s HMC 1.49 to 10.6 times, and TFP’s NUTS 0.198 to 26.6 times.

Anticorrelated draws, which hmc() can give on independent normals, have an ESS above their number, and posterior caps it at S × log10(S), or 14,400 for these 4,000 draws, with a warning each time. It capped 92 estimates in 10 runs, all on normals with greta and greta_841 (its 92 warnings are counted and silenced in Setup), and 0 of the 125 runs had a capped estimate as their smallest ESS, the number compared here.

Figure: smallest bulk ESS, not per second
Code
gg_sessions(smallest_ess_bulk, "Smallest bulk ESS") +
  scale_x_log10()

Each small dot is one session’s run, on a log scale. The large faint dot behind them is their mean.
Table: median smallest bulk ESS per second, and as a multiple of greta’s
Code
ess_per_second_ratios |>
  rename(smallest_ess_per_second = median) |>
  dt_table(digits = 2)
Table: every run
Code
efficiency |>
  select(
    session,
    model,
    method,
    smallest_ess_bulk,
    seconds,
    smallest_ess_per_second,
    largest_rhat
  ) |>
  dt_table(digits = 2)

Checks

Do the chains converge?

The largest R-hat of each run, against the 1.01 that Vehtari et al. (2021) recommend.

Code
rhat_threshold <- 1.01
runs_over_threshold <- efficiency |>
  group_by(method) |>
  summarise(over = sum(largest_rhat > rhat_threshold), runs = n())
runs_over_text <- I(paste_and(paste0(
  "`",
  runs_over_threshold$method,
  "` ",
  runs_over_threshold$over
)))
Code
gg_sessions(largest_rhat, "Largest R-hat") +
  geom_vline(xintercept = rhat_threshold, colour = "grey60")

Each small dot is one session’s run. The large faint dot behind them is their mean, and the grey line is 1.01.

From the same draws, made in benchmark.R’s bench-normals, bench-linear, bench-multiple-linear, bench-hierarchical-linear and bench-eight-schools sections.

Of each scheme’s 25 runs, this many have a largest R-hat above 1.01: greta 25, greta_841 25, windowed 15, tfp_hmc 20 and tfp_nuts 9. At 1000 warmup iterations and 1000 draws per chain, 94 of the 125 runs miss that threshold, so most of the effective samples above come from chains that have not converged by that rule.

Do the draws have the right spread?

On the normals, where the truth is known: each normal’s posterior standard deviation over its true one.

Code
posterior_scales <- posterior_summaries |>
  filter(model == "normals") |>
  mutate(
    true_sd = true_sd[as.integer(sub("x\\[([0-9]+),1\\]", "\\1", variable))],
    ratio = sd / true_sd
  )

spread_ranges <- posterior_scales |>
  group_by(method) |>
  summarise(range = format_range(ratio))
spread_text <- I(paste_and(paste0(
  "`",
  spread_ranges$method,
  "` ",
  spread_ranges$range
)))
Code
ggplot(posterior_scales, aes(x = true_sd, y = ratio, colour = method)) +
  geom_hline(yintercept = 1, colour = "grey60") +
  geom_point(alpha = 0.7) +
  scale_x_log10() +
  scale_y_log10() +
  method_scale +
  facet_wrap(~method) +
  labs(x = "True standard deviation", y = "Posterior sd / true sd")

Each dot is one normal in one session, on log scales. The grey line is a perfect estimate.

From the draws made in benchmark.R’s bench-normals section.

Each normal’s posterior standard deviation is this many times its true one: greta 0.656 to 1.11, greta_841 0.514 to 1.08, windowed 0.951 to 1.05, tfp_hmc 0.79 to 1.55 and tfp_nuts 0.96 to 1.04.

How it was measured

Code
provenance <- bind_rows(lapply(results, \(result) result$provenance))
session_numbers <- vapply(sessions, \(session) session$session, numeric(1))

# a session that pak installed from GitHub records the commit; one installed
# from a local worktree records its path, and run.R's run-versions section
# names the commit that worktree was clean at
installed_from <- vapply(
  sessions,
  \(session) session$versions[["branch"]],
  character(1)
)
installed_locally <- startsWith(installed_from, "local::")
measured_commit <- unique(sub(
  "^greta-dev/greta@",
  "",
  installed_from[!installed_locally]
))
commit_link <- I(paste0(
  "[",
  substr(measured_commit, 1, 7),
  "](https://github.com/greta-dev/greta/commit/",
  measured_commit,
  ")"
))

Every result here is from one greta commit, 92f34e3, on tf-warmup-i547, the branch of PR greta#855, which then carried the greta_841 and windowed schemes as an internal choice for comparison. Those two have since moved to windowed-adaptation-i853, a local branch built on the PR’s branch, and are not part of the PR; TFP’s two schemes are driven from benchmark.R itself. run.R’s run-versions section pins the commit: sessions 1, 2 and 3 installed it from a local worktree clean at the commit, and sessions 4 and 5 from GitHub.

run.R’s run-sessions section ran benchmark.R in 5 fresh R sessions, one after another, on Darwin 25.5.0 arm64 with 8 cores, TensorFlow 2.21.0 and TensorFlow Probability 0.25.0. Each session ran every scheme on every model, with a seed of its own from benchmark.R’s benchmark-setup section: 4 chains of 1000 warmup iterations and 1000 draws, called 2 times, as its benchmark-settings section sets. greta’s three schemes run through mcmc(), with the scheme set on greta’s internal sampler class, in its benchmark-schemes section; TFP’s two run through TFP’s windowed sampling on the same log density, in its benchmark-tfp section.

benchmark.R: benchmark-schemes
# Runs mcmc() twice on a model under each warmup scheme: the first call traces,
# the second reuses the trace. press() gives one bench_mark with a scheme
# column; next to it go the second call's draws and its tuned epsilon and
# diag_sd. The scheme is a field of greta's sampler class, read when mcmc()
# makes its sampler.
run_schemes <- function(build) {
  draws <- list()
  tuned <- list()
  marks <- press(
    scheme = bench_settings$schemes,
    {
      greta:::sampler$set("public", "adaptation", scheme, overwrite = TRUE)
      model_for_scheme <- build()
      set.seed(session_seed)
      scheme_marks <- mark(
        scheme_draws <- mcmc(
          model_for_scheme,
          warmup = bench_settings$warmup,
          n_samples = bench_settings$n_samples,
          chains = bench_settings$chains,
          n_cores = bench_settings$n_cores,
          verbose = FALSE
        ),
        iterations = bench_settings$mcmc_calls,
        check = FALSE,
        memory = FALSE,
        filter_gc = FALSE
      )
      parameters <- attr(scheme_draws, "model_info")$samplers[[1]]$parameters
      # press() evaluates this in an environment of its own, so the draws and
      # tuned values go to the lists outside it
      draws[[scheme]] <<- posterior::as_draws_array(scheme_draws)
      tuned[[scheme]] <<- list(
        epsilon = parameters$epsilon,
        diag_sd = parameters$diag_sd
      )
      scheme_marks
    }
  )
  greta:::sampler$set("public", "adaptation", "greta", overwrite = TRUE)
  list(marks = marks, draws = draws, tuned = tuned)
}
benchmark.R: benchmark-tfp
# TFP's own windowed adaptation (tfp$experimental$mcmc, as its
# windowed_adaptive_hmc() and windowed_adaptive_nuts() use it), driven from the
# model's log density with the setup windowed_adaptive_hmc() does for a joint
# distribution. Its HMC takes a fixed number of leapfrog steps. It samples the
# model's free state, which trace_values() turns into the values of the
# model's greta arrays, as mcmc() does. Each run function is traced once and
# called twice.
windowed_sampling <- reticulate::import(
  "tensorflow_probability.python.experimental.mcmc.windowed_sampling"
)
tfp <- reticulate::import("tensorflow_probability")
tf <- tensorflow::tf

tfp_windowed_run <- function(m, kind) {
  dag <- m$dag
  n_free <- length(greta:::unlist_tf(dag$example_parameters(free = TRUE)))
  float <- greta:::tf_float()
  # initial values drawn as mcmc() draws those it is not given: normal with sd
  # 0.1 on the free scale
  set.seed(session_seed)
  initial <- list(tf$constant(
    matrix(stats::rnorm(bench_settings$chains * n_free, 0, 0.1), ncol = n_free),
    dtype = float
  ))
  tensorflow::tf_function(function(seed) {
    proposal <- list(
      target_log_prob_fn = function(x) dag$tf_log_prob_function_adjusted(x),
      step_size = tf$constant(0.1, dtype = float),
      momentum_distribution = windowed_sampling$`_init_momentum`(
        initial,
        batch_shape = list(as.integer(bench_settings$chains)),
        shard_axis_names = NULL
      )
    )
    if (kind == "hmc") {
      proposal$num_leapfrog_steps <- bench_settings$tfp_leapfrog_steps
    }
    windowed_sampling$`_do_sampling`(
      kind = kind,
      proposal_kernel_kwargs = proposal,
      dual_averaging_kwargs = list(
        num_adaptation_steps = as.integer(bench_settings$warmup),
        target_accept_prob = tf$constant(0.651, dtype = float)
      ),
      num_draws = as.integer(bench_settings$n_samples),
      num_burnin_steps = as.integer(bench_settings$warmup),
      initial_position = initial,
      initial_running_variance = list(
        tfp$experimental$stats$RunningVariance$from_stats(
          num_samples = tf$zeros(list(), dtype = float),
          mean = tf$zeros_like(initial[[1]]),
          variance = tf$ones_like(initial[[1]])
        )
      ),
      bijector = tfp$bijectors$Identity(),
      trace_fn = function(...) reticulate::tuple(),
      return_final_kernel_results = FALSE,
      chain_axis_names = NULL,
      shard_axis_names = NULL,
      seed = seed
    )$all_states[[1]]
  })
}

# the free state draws, draws by chains by parameters, as a posterior draws
# array of the model's greta arrays
tfp_draws_array <- function(m, free_draws) {
  free_draws <- as.array(free_draws)
  chains <- lapply(seq_len(dim(free_draws)[2]), function(chain) {
    m$dag$trace_values(free_draws[, chain, , drop = TRUE])
  })
  values <- simplify2array(chains)
  posterior::as_draws_array(aperm(values, c(1, 3, 2)))
}

run_tfp <- function(build) {
  draws <- list()
  marks <- press(
    kind = bench_settings$tfp_kinds,
    {
      model_for_kind <- build()
      run <- tfp_windowed_run(model_for_kind, kind)
      seed <- tf$constant(as.integer(c(session_seed, 0)), dtype = tf$int32)
      kind_marks <- mark(
        free_draws <- run(seed),
        iterations = bench_settings$mcmc_calls,
        check = FALSE,
        memory = FALSE,
        filter_gc = FALSE
      )
      # press() evaluates this in an environment of its own
      draws[[kind]] <<- tfp_draws_array(model_for_kind, free_draws)
      kind_marks
    }
  )
  list(marks = marks, draws = draws)
}

Both scripts, in full:

run.R
# Runs benchmark.R against the tf-warmup-i547 branch of greta, in a fresh R
# session, in five sessions:
#
#   Rscript --quiet --vanilla posts/2026-10-07-warmup-adaptation-i853/run.R
#
# Each session writes results/session-<n>.rds, and sessions already there are
# skipped, so a rerun after a crash carries on. Each `# ---- label ----` line
# starts a section that index.qmd shows by that label.

# ---- run-versions ----
# The warmup schemes are an internal choice on the tf-warmup-i547 branch at
# this commit, added in its ancestor 59da5cc9, so it is the only version.
# Sessions 1 to 3 ran before the branch was on GitHub, so pak installed this
# commit from its worktree, clean at the commit, and their saved results record
# that local path; sessions 4 and 5 installed it from GitHub
versions <- c(
  branch = "greta-dev/greta@92f34e3003c5db24f67534d4c9abe3255ee39ccf"
)
n_sessions <- 5
post_dir <- here::here("posts", "2026-10-07-warmup-adaptation-i853")

# ---- run-sessions ----
dir.create(file.path(post_dir, "results"), showWarnings = FALSE)

for (session in seq_len(n_sessions)) {
  session_file <- file.path(
    post_dir,
    "results",
    sprintf("session-%d.rds", session)
  )
  if (file.exists(session_file)) {
    next
  }

  # cross::run_versions() installs the version into a library of its own and
  # evaluates the expression in a fresh R session against it. That session
  # cannot see this one's variables, so the script's path and the session
  # number reach it as environment variables: args_callr is passed on to
  # callr::r(), whose `env` sets variables for the child process (see ?callr::r
  # and ?callr::rcmd_safe_env, whose defaults are kept).
  measured <- cross::run_versions(
    source(Sys.getenv("GRETA_BENCH_SCRIPT"), local = TRUE)$value,
    pkgs = unname(versions),
    args_callr = list(
      env = c(
        callr::rcmd_safe_env(),
        GRETA_BENCH_SCRIPT = file.path(post_dir, "benchmark.R"),
        GRETA_BENCH_SESSION = session
      )
    )
  )

  saveRDS(
    list(
      session = session,
      versions = versions,
      results = setNames(measured$result, names(versions))
    ),
    session_file,
    compress = "xz"
  )
}
benchmark.R
# What the tf-warmup-i547 branch of greta runs, in a fresh R session: run.R
# sources it once per session, and keeps the list it ends with. Each
# `# ---- label ----` line starts a section that index.qmd shows by that label.
#
# Run it alone to try it against whichever greta is installed, which needs the
# branch's internal adaptation field:
#
#   Rscript --quiet --vanilla posts/2026-10-07-warmup-adaptation-i853/benchmark.R

# ---- benchmark-setup ----
library(greta)
library(bench)
session <- as.integer(Sys.getenv("GRETA_BENCH_SESSION", "1"))
session_seed <- 2026 - 10 - 7 + session

# ---- benchmark-settings ----
# warmup and n_samples are mcmc()'s arguments, so they count draws; mcmc_calls
# counts timed calls. The schemes are the branch's internal choices of warmup
# tuning; the TFP kinds are TFP's own windowed adaptation, for HMC and NUTS
bench_settings <- list(
  warmup = 1000,
  n_samples = 1000,
  chains = 4,
  n_cores = 4,
  mcmc_calls = 2,
  schemes = c("greta", "greta_841", "windowed"),
  tfp_kinds = c("hmc", "nuts"),
  tfp_leapfrog_steps = 8L
)

# ---- model-normals ----
# greta-dev/greta#853's test: 20 independent normals with sds from 0.01 to 100
normals_sd <- 10^seq(-2, 2, length.out = 20)
build_normals <- function() {
  x <- normal(0, normals_sd)
  model(x)
}

# ---- model-linear ----
build_linear <- function() {
  int <- normal(0, 10)
  coef <- normal(0, 10)
  sd <- cauchy(0, 3, truncation = c(0, Inf))
  mu <- int + coef * attitude$complaints
  distribution(attitude$rating) <- normal(mu, sd)
  model(int, coef, sd)
}

# ---- model-multiple-linear ----
build_multiple_linear <- function() {
  design <- as.matrix(attitude[, 2:7])
  int <- normal(0, 10)
  coefs <- normal(0, 10, dim = ncol(design))
  sd <- cauchy(0, 3, truncation = c(0, Inf))
  mu <- int + design %*% coefs
  distribution(attitude$rating) <- normal(mu, sd)
  model(int, coefs, sd)
}

# ---- model-hierarchical-linear ----
build_hierarchical_linear <- function() {
  int <- normal(0, 10)
  coef <- normal(0, 10)
  sd <- cauchy(0, 3, truncation = c(0, Inf))
  species_sd <- lognormal(0, 1)
  species_offset <- normal(0, species_sd, dim = 2)
  species_effect <- rbind(0, species_offset)
  species_id <- as.numeric(iris$Species)
  mu <- int + coef * iris$Sepal.Width + species_effect[species_id]
  distribution(iris$Sepal.Length) <- normal(mu, sd)
  model(int, coef, sd, species_sd, species_offset)
}

# ---- model-eight-schools ----
build_eight_schools <- function() {
  y <- c(28, 8, -3, 7, -1, 1, 18, 12)
  sigma_y <- c(15, 10, 16, 11, 9, 11, 10, 18)
  sigma_eta <- inverse_gamma(1, 1)
  eta <- normal(0, sigma_eta, dim = 8)
  mu_theta <- normal(0, 100)
  xi <- normal(0, 5)
  theta <- mu_theta + xi * eta
  distribution(y) <- normal(theta, sigma_y)
  model(sigma_eta, eta, mu_theta, xi)
}

# ---- benchmark-schemes ----
# Runs mcmc() twice on a model under each warmup scheme: the first call traces,
# the second reuses the trace. press() gives one bench_mark with a scheme
# column; next to it go the second call's draws and its tuned epsilon and
# diag_sd. The scheme is a field of greta's sampler class, read when mcmc()
# makes its sampler.
run_schemes <- function(build) {
  draws <- list()
  tuned <- list()
  marks <- press(
    scheme = bench_settings$schemes,
    {
      greta:::sampler$set("public", "adaptation", scheme, overwrite = TRUE)
      model_for_scheme <- build()
      set.seed(session_seed)
      scheme_marks <- mark(
        scheme_draws <- mcmc(
          model_for_scheme,
          warmup = bench_settings$warmup,
          n_samples = bench_settings$n_samples,
          chains = bench_settings$chains,
          n_cores = bench_settings$n_cores,
          verbose = FALSE
        ),
        iterations = bench_settings$mcmc_calls,
        check = FALSE,
        memory = FALSE,
        filter_gc = FALSE
      )
      parameters <- attr(scheme_draws, "model_info")$samplers[[1]]$parameters
      # press() evaluates this in an environment of its own, so the draws and
      # tuned values go to the lists outside it
      draws[[scheme]] <<- posterior::as_draws_array(scheme_draws)
      tuned[[scheme]] <<- list(
        epsilon = parameters$epsilon,
        diag_sd = parameters$diag_sd
      )
      scheme_marks
    }
  )
  greta:::sampler$set("public", "adaptation", "greta", overwrite = TRUE)
  list(marks = marks, draws = draws, tuned = tuned)
}

# ---- benchmark-tfp ----
# TFP's own windowed adaptation (tfp$experimental$mcmc, as its
# windowed_adaptive_hmc() and windowed_adaptive_nuts() use it), driven from the
# model's log density with the setup windowed_adaptive_hmc() does for a joint
# distribution. Its HMC takes a fixed number of leapfrog steps. It samples the
# model's free state, which trace_values() turns into the values of the
# model's greta arrays, as mcmc() does. Each run function is traced once and
# called twice.
windowed_sampling <- reticulate::import(
  "tensorflow_probability.python.experimental.mcmc.windowed_sampling"
)
tfp <- reticulate::import("tensorflow_probability")
tf <- tensorflow::tf

tfp_windowed_run <- function(m, kind) {
  dag <- m$dag
  n_free <- length(greta:::unlist_tf(dag$example_parameters(free = TRUE)))
  float <- greta:::tf_float()
  # initial values drawn as mcmc() draws those it is not given: normal with sd
  # 0.1 on the free scale
  set.seed(session_seed)
  initial <- list(tf$constant(
    matrix(stats::rnorm(bench_settings$chains * n_free, 0, 0.1), ncol = n_free),
    dtype = float
  ))
  tensorflow::tf_function(function(seed) {
    proposal <- list(
      target_log_prob_fn = function(x) dag$tf_log_prob_function_adjusted(x),
      step_size = tf$constant(0.1, dtype = float),
      momentum_distribution = windowed_sampling$`_init_momentum`(
        initial,
        batch_shape = list(as.integer(bench_settings$chains)),
        shard_axis_names = NULL
      )
    )
    if (kind == "hmc") {
      proposal$num_leapfrog_steps <- bench_settings$tfp_leapfrog_steps
    }
    windowed_sampling$`_do_sampling`(
      kind = kind,
      proposal_kernel_kwargs = proposal,
      dual_averaging_kwargs = list(
        num_adaptation_steps = as.integer(bench_settings$warmup),
        target_accept_prob = tf$constant(0.651, dtype = float)
      ),
      num_draws = as.integer(bench_settings$n_samples),
      num_burnin_steps = as.integer(bench_settings$warmup),
      initial_position = initial,
      initial_running_variance = list(
        tfp$experimental$stats$RunningVariance$from_stats(
          num_samples = tf$zeros(list(), dtype = float),
          mean = tf$zeros_like(initial[[1]]),
          variance = tf$ones_like(initial[[1]])
        )
      ),
      bijector = tfp$bijectors$Identity(),
      trace_fn = function(...) reticulate::tuple(),
      return_final_kernel_results = FALSE,
      chain_axis_names = NULL,
      shard_axis_names = NULL,
      seed = seed
    )$all_states[[1]]
  })
}

# the free state draws, draws by chains by parameters, as a posterior draws
# array of the model's greta arrays
tfp_draws_array <- function(m, free_draws) {
  free_draws <- as.array(free_draws)
  chains <- lapply(seq_len(dim(free_draws)[2]), function(chain) {
    m$dag$trace_values(free_draws[, chain, , drop = TRUE])
  })
  values <- simplify2array(chains)
  posterior::as_draws_array(aperm(values, c(1, 3, 2)))
}

run_tfp <- function(build) {
  draws <- list()
  marks <- press(
    kind = bench_settings$tfp_kinds,
    {
      model_for_kind <- build()
      run <- tfp_windowed_run(model_for_kind, kind)
      seed <- tf$constant(as.integer(c(session_seed, 0)), dtype = tf$int32)
      kind_marks <- mark(
        free_draws <- run(seed),
        iterations = bench_settings$mcmc_calls,
        check = FALSE,
        memory = FALSE,
        filter_gc = FALSE
      )
      # press() evaluates this in an environment of its own
      draws[[kind]] <<- tfp_draws_array(model_for_kind, free_draws)
      kind_marks
    }
  )
  list(marks = marks, draws = draws)
}

# ---- bench-normals ----
normals_schemes <- run_schemes(build_normals)
normals_tfp <- run_tfp(build_normals)

# ---- bench-linear ----
linear_schemes <- run_schemes(build_linear)
linear_tfp <- run_tfp(build_linear)

# ---- bench-multiple-linear ----
multiple_linear_schemes <- run_schemes(build_multiple_linear)
multiple_linear_tfp <- run_tfp(build_multiple_linear)

# ---- bench-hierarchical-linear ----
hierarchical_linear_schemes <- run_schemes(build_hierarchical_linear)
hierarchical_linear_tfp <- run_tfp(build_hierarchical_linear)

# ---- bench-eight-schools ----
eight_schools_schemes <- run_schemes(build_eight_schools)
eight_schools_tfp <- run_tfp(build_eight_schools)

# ---- benchmark-provenance ----
session_info <- sessioninfo::session_info()
# greta_sitrep() reports through messages, so capture those
greta_sitrep_report <- utils::capture.output(greta_sitrep(), type = "message")

provenance <- data.frame(
  greta_version = as.character(packageVersion("greta")),
  r_version = paste(R.version$major, R.version$minor, sep = "."),
  tensorflow_version = tensorflow::tf$version$VERSION,
  tfp_version = tfp$`__version__`,
  machine = paste(
    Sys.info()[c("sysname", "release", "machine")],
    collapse = " "
  ),
  cores = parallel::detectCores()
)

# ---- benchmark-results ----
# the value source() returns to run.R: every bench_mark object as press() and
# mark() made it, with each run's draws and tuned values beside them
list(
  session_seed = session_seed,
  bench_settings = bench_settings,
  normals_sd = normals_sd,
  provenance = provenance,
  session_info = session_info,
  greta_sitrep_report = greta_sitrep_report,
  schemes = list(
    normals = normals_schemes,
    linear = linear_schemes,
    multiple_linear = multiple_linear_schemes,
    hierarchical_linear = hierarchical_linear_schemes,
    eight_schools = eight_schools_schemes
  ),
  tfp = list(
    normals = normals_tfp,
    linear = linear_tfp,
    multiple_linear = multiple_linear_tfp,
    hierarchical_linear = hierarchical_linear_tfp,
    eight_schools = eight_schools_tfp
  )
)

The models

Twenty independent normals with standard deviations from 0.01 to 100, greta#853’s test: their scales span four orders of magnitude, so a diag_sd that is off on one scale shows. The other four are greta’s examples, written out in maths in the PR’s main post.

benchmark.R: model-normals
# greta-dev/greta#853's test: 20 independent normals with sds from 0.01 to 100
normals_sd <- 10^seq(-2, 2, length.out = 20)
build_normals <- function() {
  x <- normal(0, normals_sd)
  model(x)
}
benchmark.R: model-linear
build_linear <- function() {
  int <- normal(0, 10)
  coef <- normal(0, 10)
  sd <- cauchy(0, 3, truncation = c(0, Inf))
  mu <- int + coef * attitude$complaints
  distribution(attitude$rating) <- normal(mu, sd)
  model(int, coef, sd)
}
benchmark.R: model-multiple-linear
build_multiple_linear <- function() {
  design <- as.matrix(attitude[, 2:7])
  int <- normal(0, 10)
  coefs <- normal(0, 10, dim = ncol(design))
  sd <- cauchy(0, 3, truncation = c(0, Inf))
  mu <- int + design %*% coefs
  distribution(attitude$rating) <- normal(mu, sd)
  model(int, coefs, sd)
}
benchmark.R: model-hierarchical-linear
build_hierarchical_linear <- function() {
  int <- normal(0, 10)
  coef <- normal(0, 10)
  sd <- cauchy(0, 3, truncation = c(0, Inf))
  species_sd <- lognormal(0, 1)
  species_offset <- normal(0, species_sd, dim = 2)
  species_effect <- rbind(0, species_offset)
  species_id <- as.numeric(iris$Species)
  mu <- int + coef * iris$Sepal.Width + species_effect[species_id]
  distribution(iris$Sepal.Length) <- normal(mu, sd)
  model(int, coef, sd, species_sd, species_offset)
}
benchmark.R: model-eight-schools
build_eight_schools <- function() {
  y <- c(28, 8, -3, 7, -1, 1, 18, 12)
  sigma_y <- c(15, 10, 16, 11, 9, 11, 10, 18)
  sigma_eta <- inverse_gamma(1, 1)
  eta <- normal(0, sigma_eta, dim = 8)
  mu_theta <- normal(0, 100)
  xi <- normal(0, 5)
  theta <- mu_theta + xi * eta
  distribution(y) <- normal(theta, sigma_y)
  model(sigma_eta, eta, mu_theta, xi)
}

Session details

What the branch reported about itself in session 1, from benchmark.R’s benchmark-provenance section.

greta_sitrep()
Code
writeLines(results[[1]]$greta_sitrep_report)
ℹ checking if python available
✔ python (v3.12) available

ℹ checking if TensorFlow available
✔ TensorFlow (v2.21.0) available

ℹ checking if TensorFlow Probability available
✔ TensorFlow Probability (v0.25.0) available

ℹ greta conda environment: not used (managed (uv) environment active)
• backend: "managed (uv) environment"
• selected via: default
✔ offline-ready: `UV_OFFLINE`=1 is set and the uv cache is populated
ℹ All dependencies available; greta is ready to use!
sessioninfo::session_info()
Code
# a saved session_info() prints as a plain list unless sessioninfo's print
# method is registered, which loading its namespace does
invisible(loadNamespace("sessioninfo"))
results[[1]]$session_info
─ Session info ───────────────────────────────────────────────────────────────
 setting  value
 version  R version 4.6.1 (2026-06-24)
 os       macOS Tahoe 26.5.2
 system   aarch64, darwin23
 ui       X11
 language (EN)
 collate  en_AU.UTF-8
 ctype    en_AU.UTF-8
 tz       Australia/Hobart
 date     2026-10-07
 pandoc   3.10 @ /opt/homebrew/bin/pandoc
 quarto   1.9.36 @ /usr/local/bin/quarto

─ Packages ───────────────────────────────────────────────────────────────────
 package        * version    date (UTC) lib source
 abind            1.4-8      2024-09-12 [1] CRAN (R 4.6.0)
 backports        1.5.1      2026-04-03 [1] CRAN (R 4.6.0)
 base64enc        0.1-6      2026-02-02 [1] CRAN (R 4.6.0)
 bench          * 1.1.4      2025-01-16 [3] CRAN (R 4.6.0)
 callr            3.8.0      2026-06-05 [1] CRAN (R 4.6.0)
 checkmate        2.3.4      2026-02-03 [2] CRAN (R 4.6.0)
 cli              3.6.6      2026-04-09 [1] CRAN (R 4.6.0)
 coda             0.19-4.1   2024-01-31 [1] CRAN (R 4.6.0)
 codetools        0.2-20     2024-03-31 [3] CRAN (R 4.6.1)
 crayon           1.5.3      2024-06-20 [1] CRAN (R 4.6.0)
 digest           0.6.39     2025-11-19 [1] CRAN (R 4.6.0)
 distributional   0.9.0      2026-09-08 [2] CRAN (R 4.6.1)
 future           1.76.0     2026-09-24 [1] CRAN (R 4.6.1)
 generics         0.1.4      2025-05-09 [2] CRAN (R 4.6.0)
 globals          0.19.1     2026-03-13 [1] CRAN (R 4.6.0)
 glue             1.8.1      2026-04-17 [1] CRAN (R 4.6.0)
 greta          * 0.6.0.9000 2026-10-06 [1] local
 hms              1.1.4      2025-10-17 [1] CRAN (R 4.6.0)
 jsonlite         2.0.0      2025-03-27 [1] CRAN (R 4.6.0)
 lattice          0.23-1     2026-08-12 [3] CRAN (R 4.6.1)
 lifecycle        1.0.5      2026-01-08 [1] CRAN (R 4.6.0)
 listenv          1.1.0      2026-10-04 [1] CRAN (R 4.6.1)
 magrittr         2.0.5      2026-04-04 [1] CRAN (R 4.6.0)
 Matrix           1.7-6      2026-07-25 [3] CRAN (R 4.6.1)
 parallelly       1.48.0     2026-06-29 [1] CRAN (R 4.6.1)
 pillar           1.11.1     2025-09-17 [2] CRAN (R 4.6.0)
 pkgconfig        2.0.3      2019-09-22 [1] CRAN (R 4.6.0)
 png              0.1-9      2026-03-15 [1] CRAN (R 4.6.0)
 posterior        1.7.0      2026-04-01 [2] CRAN (R 4.6.0)
 prettyunits      1.2.0      2023-09-24 [1] CRAN (R 4.6.0)
 processx         3.9.0      2026-04-22 [1] CRAN (R 4.6.0)
 progress         1.2.3      2023-12-06 [1] CRAN (R 4.6.0)
 R6               2.6.1      2025-02-15 [1] CRAN (R 4.6.0)
 Rcpp             1.1.2      2026-07-05 [1] CRAN (R 4.6.1)
 reticulate       1.47.0     2026-09-03 [1] CRAN (R 4.6.1)
 rlang            1.3.0      2026-07-05 [1] CRAN (R 4.6.1)
 sessioninfo      1.2.4      2026-06-04 [3] CRAN (R 4.6.0)
 tensorA          0.36.2.1   2023-12-13 [2] CRAN (R 4.6.0)
 tensorflow       2.20.0     2025-08-22 [1] CRAN (R 4.6.0)
 tfautograph      0.3.2      2021-09-17 [1] CRAN (R 4.6.0)
 tfruns           1.5.4      2025-08-20 [1] CRAN (R 4.6.0)
 tibble           3.3.1      2026-01-11 [2] CRAN (R 4.6.0)
 utf8             1.2.6      2025-06-08 [2] CRAN (R 4.6.0)
 vctrs            0.7.3      2026-04-11 [1] CRAN (R 4.6.0)
 whisker          0.4.1      2022-12-05 [1] CRAN (R 4.6.0)
 withr            3.0.3      2026-06-19 [1] CRAN (R 4.6.0)

 [1] /private/var/folders/cz/hrp33k3s0jlcb_gg9snlqh580000gq/T/RtmploxQXB/file9e8f2f802efc
 [2] /Users/nick_1/Library/R/arm64/4.6/library
 [3] /Library/Frameworks/R.framework/Versions/4.6/Resources/library
 * ── Packages attached to the search path.

─ Python configuration ───────────────────────────────────────────────────────
 /Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF/bin/python
 /Users/nick_1/.local/share/uv/python/cpython-3.12.13-macos-aarch64-none/lib/lib/libpython3.12.dylib
 /Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF:/Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF
 /Users/nick_1/.local/share/uv/python/cpython-3.12.13-macos-aarch64-none/lib/python312.zip:/Users/nick_1/.local/share/uv/python/cpython-3.12.13-macos-aarch64-none/lib/python3.12:/Users/nick_1/.local/share/uv/python/cpython-3.12.13-macos-aarch64-none/lib/python3.12/lib-dynload:/Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF/lib/python3.12/site-packages
 /Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF
 /Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF
 /Users/nick_1/.local/share/uv/python/cpython-3.12.13-macos-aarch64-none
 /Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF
 /Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF/bin/activate_this.py
 /Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF/bin/python
 /Users/nick_1/.local/share/uv/python/cpython-3.12.13-macos-aarch64-none/bin/python3.12
 3.12.13 (main, Mar 20 2026, 00:20:20) [Clang 22.1.1 ]
 3, 12
 64bit
 FALSE
 False
 /Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF/lib/python3.12/site-packages/numpy, 2                                                                                          , 5                                                                                          , 3                                                                                          
 tensorflow
 /Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF/lib/python3.12/site-packages/tensorflow
 TRUE
 /Users/nick_1/.cache/uv/archive-v0/L7vrsM8H3lbjoC0FIwGSF/bin/python
 py_require()
 TRUE

──────────────────────────────────────────────────────────────────────────────

References

Girolami, Mark, and Ben Calderhead. 2011. “Riemann Manifold Langevin and Hamiltonian Monte Carlo Methods.” Journal of the Royal Statistical Society: Series B (Statistical Methodology) 73 (2): 123–214. https://doi.org/10.1111/j.1467-9868.2010.00765.x.
Hoffman, Matthew D., and Andrew Gelman. 2014. “The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo.” Journal of Machine Learning Research 15 (47): 1593–623. https://jmlr.org/papers/v15/hoffman14a.html.
Vehtari, Aki, Andrew Gelman, Daniel Simpson, Bob Carpenter, and Paul-Christian Bürkner. 2021. “Rank-Normalization, Folding, and Localization: An Improved R̂ for Assessing Convergence of MCMC (with Discussion).” Bayesian Analysis 16 (2): 667–718. https://doi.org/10.1214/20-BA1221.