greta#855 (tf-warmup-i547): TFP’s loops, tried and not adopted

TensorFlow Probability’s sample_chain() and minimize() timed against greta’s own sampling loop and opt(), on the PR’s branch and two models. Once traced, sample_chain() takes as long as greta’s loop, and minimize() takes longer than opt(), so the PR keeps greta’s own loops.

suppressPackageStartupMessages({
  library(dplyr)
  library(tidyr)
  library(ggplot2)
  library(bench)
  library(here)
})

post_path <- "posts/2026-10-07-tfp-alternatives-i547"

# 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
  )
}
# Quarto escapes the markdown in an inline `{r}` result unless it is wrapped
# in I(), so both link helpers return I()
file_link <- function(file) {
  I(paste0("[`", file, "`](", github_file_url(file), ")"))
}

# a link to the lines of a script's `# ---- label ----` section
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, ")"
  ))
}

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
provenance <- results[[1]]$provenance
python_version <- sub(
  " .*",
  "",
  results[[1]]$session_info$python$version_string
)
model_levels <- names(results[[1]]$sampling)

loop_levels <- c("greta", "tfp")
loop_colours <- c(greta = "#009E73", tfp = "#E69F00")

# a loop's first timed call on a model, apart from the calls after it
label_calls <- function(call) {
  factor(
    ifelse(call == 1, "first call", "later calls"),
    levels = c("first call", "later calls")
  )
}

# one row per timed call: bench keeps every call's time in the `time` column
every_call <- function(marks, ...) {
  marks |>
    as_tibble() |>
    mutate(loop = factor(as.character(expression), levels = loop_levels)) |>
    select(loop, ..., time) |>
    mutate(call = lapply(time, seq_along)) |>
    unnest(c(time, call)) |>
    mutate(seconds = as.numeric(time), calls = label_calls(call), .keep = "unused")
}

# every timed call, from every session and model
sampling_calls <- lapply(seq_along(results), function(i) {
  lapply(model_levels, function(model_name) {
    every_call(results[[i]]$sampling[[model_name]]$marks) |>
      mutate(session = sessions[[i]]$session, model = model_name)
  }) |>
    bind_rows()
}) |>
  bind_rows() |>
  mutate(model = factor(model, levels = model_levels))

opt_calls <- lapply(seq_along(results), function(i) {
  lapply(model_levels, function(model_name) {
    every_call(results[[i]]$opt[[model_name]], max_iterations) |>
      mutate(session = sessions[[i]]$session, model = model_name)
  }) |>
    bind_rows()
}) |>
  bind_rows() |>
  mutate(
    model = factor(model, levels = model_levels),
    iterations = factor(paste(max_iterations, "iterations"))
  )

# each loop's median call, and how many times as long TFP's loop takes as
# greta's
loop_medians <- function(calls, ...) {
  calls |>
    group_by(..., loop) |>
    summarise(seconds = median(seconds), .groups = "drop") |>
    pivot_wider(names_from = loop, values_from = seconds) |>
    mutate(tfp_over_greta = tfp / greta)
}
sampling_medians <- loop_medians(sampling_calls, model, calls)
sampling_ratios <- split(
  sampling_medians$tfp_over_greta,
  sampling_medians$calls
)
opt_medians <- loop_medians(opt_calls, model, max_iterations)

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

# "a to b", or one number when the two round the same
format_range <- function(x) {
  ends <- format_number(range(x))
  if (ends[1] == ends[2]) {
    return(ends[1])
  }
  paste(ends, collapse = " to ")
}

# A raincloud for each loop: the density of its calls above, their box plot
# in the middle, and every call as a dot below. The summaries are drawn first
# and faint, so the dots sit on top.
raincloud <- list(
  ggdist::stat_halfeye(
    aes(fill = loop),
    width = 0.4,
    justification = -0.3,
    .width = 0,
    point_colour = NA,
    alpha = 0.3
  ),
  geom_boxplot(aes(fill = loop), width = 0.25, outlier.shape = NA, alpha = 0.3),
  # a fixed dot size, so a call looks the same size in every panel
  ggdist::stat_dots(
    aes(fill = loop),
    side = "bottom",
    layout = "swarm",
    justification = 1.3,
    scale = 0.75,
    binwidth = unit(2.6, "mm"),
    overflow = "compress",
    colour = NA,
    alpha = 0.6
  ),
  scale_colour_manual(values = loop_colours, guide = "none"),
  scale_fill_manual(values = loop_colours, guide = "none")
)
raincloud_caption <- "Each dot is one call, from every session. Above the dots, the curve is the density of the calls and the box marks their median and quartiles."

# 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)
}

greta#855 runs mcmc()’s sampling and opt()’s optimiser in loops of greta’s own, inside TensorFlow; the main post times them. This post times TensorFlow Probability’s loops for the same jobs, sample_chain() and minimize(), against greta’s, on the PR’s branch and two models. Once traced, sample_chain() takes 0.999 times as long as greta’s loop, and minimize() takes 1.41 to 1.59 times as long as opt(). Neither is faster, so the PR keeps greta’s own loops.

How long does sample_chain() take?

Each session ran hmc() on each model for 2000 iterations on 4 chains, 6 times through each loop, from the state one mcmc() call had tuned.

ggplot(sampling_calls, aes(x = seconds, y = loop, colour = loop)) +
  raincloud +
  scale_y_discrete(labels = c(greta = "greta's loop", tfp = "sample_chain()")) +
  facet_wrap(vars(model, calls), ncol = 2, scales = "free_x") +
  labs(x = "Seconds", y = NULL)

Each dot is one call, from every session. Above the dots, the curve is the density of the calls and the box marks their median and quartiles.

Timed by benchmark.R’s benchmark-sampling section, with the kernel from its tfp-leapfrog-kernel section.

Once traced, sample_chain() takes 0.999 times as long as greta’s loop. Its first call, which traces it, takes 1.53 to 1.54 times as long as greta’s first call, which reuses the trace of the mcmc() call that tuned the sampler.

Table: median seconds of each loop’s calls

dt_table(sampling_medians, digits = 3)

How long does minimize() take?

Each session ran opt() with adam(), and minimize() with the same Keras optimiser, 6 times each on each model at 100 and 1000 iterations, building and tracing the loop afresh on every call, as opt() does.

# what the extra iterations cost opt(): the difference between its median call
# at the most and at the fewest iterations
iteration_cost <- opt_medians |>
  group_by(model) |>
  summarise(
    fewest = greta[which.min(max_iterations)],
    extra = greta[which.max(max_iterations)] - fewest,
    .groups = "drop"
  )

ggplot(opt_calls, aes(x = seconds, y = loop, colour = loop)) +
  raincloud +
  scale_y_discrete(labels = c(greta = "opt()", tfp = "minimize()")) +
  facet_wrap(vars(model, iterations), ncol = 2, scales = "free_x") +
  labs(x = "Seconds", y = NULL)

Each dot is one call, from every session. Above the dots, the curve is the density of the calls and the box marks their median and quartiles.

Timed by benchmark.R’s benchmark-opt section.

minimize() takes 1.41 to 1.59 times as long as opt(). Most of an opt() call is building and tracing its loop: on the two models, its median call takes 0.0898 to 0.132 seconds at 100 iterations, and 900 more iterations add 0.0178 to 0.0332 seconds.

Table: median seconds of each loop’s calls

dt_table(opt_medians, digits = 3)

Do both loops sample the same posterior?

For each free parameter in each session, the difference between the two loops’ posterior means, in units of its Monte Carlo standard error (Vehtari et al. 2021): a z beyond about ±3 would say they sample different posteriors.

loop_means <- lapply(seq_along(results), function(i) {
  lapply(model_levels, function(model_name) {
    draws <- results[[i]]$sampling[[model_name]]$draws
    lapply(loop_levels, function(loop) {
      tibble(
        session = sessions[[i]]$session,
        model = model_name,
        loop = loop,
        posterior::summarise_draws(draws[[loop]], "mean", "mcse_mean")
      )
    }) |>
      bind_rows()
  }) |>
    bind_rows()
}) |>
  bind_rows()

agreement <- loop_means |>
  pivot_wider(
    id_cols = c(session, model, variable),
    names_from = loop,
    values_from = c(mean, mcse_mean)
  ) |>
  mutate(
    model = factor(model, levels = model_levels),
    # the draws are of the free state, whose columns have no names, so
    # posterior names them ...1, ...2 and so on
    variable = sub("^[.]+", "free parameter ", variable),
    z = (mean_greta - mean_tfp) / sqrt(mcse_mean_greta^2 + mcse_mean_tfp^2)
  )

ggplot(agreement, aes(x = z, y = variable)) +
  annotate("rect", xmin = -3, xmax = 3, ymin = -Inf, ymax = Inf, alpha = 0.15) +
  geom_point(alpha = 0.6, size = 2) +
  facet_wrap(vars(model), scales = "free_y") +
  expand_limits(x = c(-4, 4)) +
  labs(x = "z", y = NULL)

Each dot is one free parameter in one session. The grey band marks z between -3 and 3.

Drawn by benchmark.R’s benchmark-sampling section, in each loop’s last call.

The largest |z| of the 45 comparisons is 2.48, so none says the loops sample different posteriors.

Table: every comparison

dt_table(agreement, digits = 3)

How it was measured

run.R ran benchmark.R top to bottom in 5 fresh R sessions, with the settings in its benchmark-settings section. Each session ran one version of greta, the PR’s branch at commit 5bdab31e, which is not on GitHub, so pak installed it from a local worktree, as run-versions says. Both loops of each pair ran in the same session, greta’s calls before TFP’s. The sessions ran on Python 3.12.13, TensorFlow 2.21.0, TensorFlow Probability 0.25.0 and reticulate 1.47.0, on Darwin 25.5.0 arm64 with 8 cores.

# Runs benchmark.R against the tf-sampler-i547 branch of greta, in a fresh R
# session, in five sessions:
#
#   Rscript --quiet --vanilla posts/2026-10-07-tfp-alternatives-i547/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 ----
# Both ways of running each loop are measured in the same session of the same
# greta, so the branch is the only version. It was not on GitHub when this
# ran, so pak installed this commit from its worktree, clean at the commit, and
# the saved results record that local path
versions <- c(
  branch = "greta-dev/greta@5bdab31e9ec9295c845a4fac9d1201e1aabb9c93"
)
n_sessions <- 5
post_dir <- here::here("posts", "2026-10-07-tfp-alternatives-i547")

# ---- 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"
  )
}

# What the tf-sampler-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.
#
# It reaches into greta's internals, so it needs the branch:
#
#   Rscript --quiet --vanilla posts/2026-10-07-tfp-alternatives-i547/benchmark.R

# ---- benchmark-setup ----
library(greta)
library(bench)
session <- as.integer(Sys.getenv("GRETA_BENCH_SESSION", "1"))
session_seed <- 2026 - 10 - 7 + session
tfp <- reticulate::import("tensorflow_probability")
tf <- tensorflow::tf
builtins <- reticulate::import_builtins()

# ---- benchmark-settings ----
# sampling_iterations is the length of each timed sampling call;
# opt_max_iterations are the iteration limits opt() is timed at, with a
# tolerance of 0 so it runs all of them; the *_calls settings count timed calls
bench_settings <- list(
  chains = 4,
  warmup = 1000,
  sampling_iterations = 2000L,
  sampling_calls = 6,
  opt_max_iterations = c(100L, 1000L),
  opt_learning_rate = 0.1,
  opt_calls = 6
)

# ---- 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-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)
}

# ---- tfp-leapfrog-kernel ----
# hmc() as a TFP kernel. TFP's HamiltonianMonteCarlo takes one leapfrog count
# when it is built, and reads it from its kernel results instead when built
# with store_parameters_in_results = TRUE; this wrapper writes a fresh count,
# between Lmin and Lmax, into those results every step. MetropolisHastings
# passes accepted_results to the uncalibrated kernel, so the count goes there.
RandomLeapfrog <- reticulate::PyClass(
  "RandomLeapfrog",
  inherit = tfp$mcmc$TransitionKernel,
  defs = list(
    `__init__` = function(self, inner_kernel, l_min, l_max) {
      self$inner_kernel <- inner_kernel
      self$l_min <- l_min
      self$l_max <- l_max
      NULL
    },
    is_calibrated = builtins$property(function(self) TRUE),
    bootstrap_results = function(self, init_state) {
      self$inner_kernel$bootstrap_results(init_state)
    },
    one_step = function(self, current_state, previous_kernel_results, seed = NULL) {
      seeds <- tf$unstack(tfp$random$split_seed(seed, n = 2L))
      n_leapfrog <- tf$random$stateless_uniform(
        shape = list(),
        seed = seeds[[1]],
        minval = self$l_min,
        maxval = self$l_max + 1L,
        dtype = tf$int32
      )
      accepted <- previous_kernel_results$accepted_results$`_replace`(
        num_leapfrog_steps = n_leapfrog
      )
      results <- previous_kernel_results$`_replace`(accepted_results = accepted)
      self$inner_kernel$one_step(current_state, results, seed = seeds[[2]])
    }
  )
)

# ---- benchmark-sampling ----
# A model's sampler, tuned by one mcmc() call, then run for
# sampling_iterations two ways from the same state, step sizes and leapfrog
# range: through greta's own traced loop, called as mcmc() calls it, and
# through tfp$mcmc$sample_chain() with the wrapper kernel. Each is traced on
# its first call and reuses the trace after, so every call is kept and the
# post separates the first. Each also returns its draws, for the posterior
# means to compare.
bench_sampling <- function(m) {
  set.seed(session_seed)
  tuned <- mcmc(
    m,
    warmup = bench_settings$warmup,
    n_samples = 10,
    chains = bench_settings$chains,
    verbose = FALSE
  )
  sampler <- attr(tuned, "model_info")$samplers[[1]]
  dag <- m$dag
  float <- greta:::tf_float()
  n_free <- sampler$n_free
  param_vec <- unlist(sampler$sampler_parameter_values())
  diag_sd <- param_vec[3 + seq_len(n_free)]
  step_sizes <- param_vec[3] * diag_sd / sum(diag_sd)
  initial <- tensorflow::as_tensor(sampler$free_state, dtype = float)
  n <- bench_settings$sampling_iterations

  greta_loop <- function() {
    sampler$tf_iterations(
      free_state = initial,
      n_iterations = tensorflow::as_tensor(n),
      first_iteration = tensorflow::as_tensor(10000L),
      warmup = tensorflow::as_tensor(0L),
      sampling_start = tensorflow::as_tensor(10000L),
      thin = tensorflow::as_tensor(1L),
      param_vec = tensorflow::as_tensor(
        param_vec,
        dtype = float,
        shape = length(param_vec)
      ),
      tuning = tensorflow::as_tensor(rep(0, 6), dtype = float),
      welford_mean = tensorflow::as_tensor(rep(0, n_free), dtype = float, shape = n_free),
      welford_m2 = tensorflow::as_tensor(rep(0, n_free), dtype = float, shape = n_free),
      seed = tensorflow::as_tensor(c(session_seed, 0L), dtype = tf$int32)
    )$draws
  }

  tfp_chain <- tensorflow::tf_function(
    function(state) {
      inner <- tfp$mcmc$HamiltonianMonteCarlo(
        target_log_prob_fn = dag$tf_log_prob_function_adjusted,
        step_size = tf$constant(step_sizes, dtype = tf$float64),
        num_leapfrog_steps = as.integer(param_vec[1]),
        store_parameters_in_results = TRUE
      )
      kernel <- RandomLeapfrog(
        inner,
        l_min = as.integer(param_vec[1]),
        l_max = as.integer(param_vec[2])
      )
      tfp$mcmc$sample_chain(
        num_results = n,
        current_state = state,
        kernel = kernel,
        trace_fn = function(current_state, kernel_results) {
          kernel_results$log_accept_ratio
        },
        seed = as.integer(c(session_seed, 1L))
      )$all_states
    },
    input_signature = list(tf$TensorSpec(
      shape = list(as.integer(bench_settings$chains), n_free),
      dtype = float
    ))
  )
  tfp_loop <- function() tfp_chain(initial)

  draws <- list(greta = NULL, tfp = NULL)
  marks <- mark(
    greta = draws$greta <- as.array(greta_loop()),
    tfp = draws$tfp <- as.array(tfp_loop()),
    iterations = bench_settings$sampling_calls,
    check = FALSE,
    memory = FALSE,
    filter_gc = FALSE
  )
  list(marks = marks, draws = draws)
}

# ---- benchmark-opt ----
# opt() with adam(), against tfp$math$minimize() running the same Keras
# optimiser with greta's convergence rule: stop once successive objectives
# differ by at most the tolerance. That rule is a ConvergenceCriterion
# subclass written in R. opt() starts from small random values on the free
# scale, as it does by default, and minimize() from zero. Both build and trace
# their loop afresh on every call, as opt() does.
AbsoluteChange <- reticulate::PyClass(
  "AbsoluteChange",
  inherit = tfp$optimizer$convergence_criteria$ConvergenceCriterion,
  defs = list(
    `__init__` = function(self, tolerance) {
      tfp$optimizer$convergence_criteria$ConvergenceCriterion$`__init__`(
        self,
        min_num_steps = 1L,
        name = "AbsoluteChange"
      )
      self$tolerance <- tolerance
      NULL
    },
    `_bootstrap` = function(self, loss, grads, parameters) {
      loss
    },
    `_one_step` = function(self, step, loss, grads, parameters, auxiliary_state) {
      reticulate::tuple(
        tf$less_equal(tf$abs(loss - auxiliary_state), self$tolerance),
        loss
      )
    }
  )
)

opt_tfp <- function(m, max_iterations, tolerance) {
  dag <- m$dag
  float <- greta:::tf_float()
  n_free <- length(greta:::unlist_tf(dag$example_parameters(free = TRUE)))
  free_state <- tf$Variable(matrix(0, 1, n_free), dtype = float)
  optimiser <- tf$keras$optimizers$Adam(
    learning_rate = bench_settings$opt_learning_rate
  )
  optimiser$build(list(free_state))
  # minimize() calls the loss with a seed, and falls back to calling it
  # without one only on Python's TypeError, which an R error is not
  loss <- function(seed = NULL) {
    tf$reshape(-dag$one_row_log_prob(free_state)$adjusted, shape = list())
  }
  run <- tensorflow::tf_function(function() {
    tfp$math$minimize(
      loss_fn = loss,
      num_steps = max_iterations,
      optimizer = optimiser,
      convergence_criterion = AbsoluteChange(tolerance = tolerance),
      trainable_variables = list(free_state)
    )
  })
  run()
  free_state$numpy()
}

bench_opt <- function(m) {
  bench::press(
    max_iterations = bench_settings$opt_max_iterations,
    {
      set.seed(session_seed)
      mark(
        greta = opt(
          m,
          optimiser = adam(learning_rate = bench_settings$opt_learning_rate),
          max_iterations = max_iterations,
          tolerance = 0
        ),
        tfp = opt_tfp(m, max_iterations, tolerance = 0),
        iterations = bench_settings$opt_calls,
        check = FALSE,
        memory = FALSE,
        filter_gc = FALSE
      )
    }
  )
}

# ---- bench-linear ----
linear_model <- build_linear()
linear_sampling <- bench_sampling(linear_model)
linear_opt <- bench_opt(linear_model)

# ---- bench-hierarchical-linear ----
hierarchical_linear_model <- build_hierarchical_linear()
hierarchical_linear_sampling <- bench_sampling(hierarchical_linear_model)
hierarchical_linear_opt <- bench_opt(hierarchical_linear_model)

# ---- 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 = tf$version$VERSION,
  tfp_version = tfp$`__version__`,
  reticulate_version = as.character(packageVersion("reticulate")),
  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 mark() and
# press() made it, with each sampling run's draws beside it
list(
  session_seed = session_seed,
  bench_settings = bench_settings,
  provenance = provenance,
  session_info = session_info,
  greta_sitrep_report = greta_sitrep_report,
  sampling = list(
    linear = linear_sampling,
    hierarchical_linear = hierarchical_linear_sampling
  ),
  opt = list(
    linear = linear_opt,
    hierarchical_linear = hierarchical_linear_opt
  )
)

The TFP loops

Neither TFP loop fits greta as it stands, so each needed a subclass of a TFP class, written in R with reticulate::PyClass():

• sample_chain() steps one kernel, built before its loop, so hmc()’s number of leapfrog steps would be fixed for the whole call, where greta draws it afresh every iteration. TFP’s HamiltonianMonteCarlo reads the number from its kernel results when built with store_parameters_in_results = TRUE, so a wrapper kernel, RandomLeapfrog, writes a fresh number there every step. Both loops start from the same state, with the same step sizes and range of leapfrog steps.
• minimize() stops when a convergence criterion says so, and none of TFP’s is greta’s: stop once two successive objectives differ by at most tolerance. The AbsoluteChange criterion does that. The tolerance here is 0, so both loops run every iteration. opt() starts from small random values and minimize() from zero.

# hmc() as a TFP kernel. TFP's HamiltonianMonteCarlo takes one leapfrog count
# when it is built, and reads it from its kernel results instead when built
# with store_parameters_in_results = TRUE; this wrapper writes a fresh count,
# between Lmin and Lmax, into those results every step. MetropolisHastings
# passes accepted_results to the uncalibrated kernel, so the count goes there.
RandomLeapfrog <- reticulate::PyClass(
  "RandomLeapfrog",
  inherit = tfp$mcmc$TransitionKernel,
  defs = list(
    `__init__` = function(self, inner_kernel, l_min, l_max) {
      self$inner_kernel <- inner_kernel
      self$l_min <- l_min
      self$l_max <- l_max
      NULL
    },
    is_calibrated = builtins$property(function(self) TRUE),
    bootstrap_results = function(self, init_state) {
      self$inner_kernel$bootstrap_results(init_state)
    },
    one_step = function(self, current_state, previous_kernel_results, seed = NULL) {
      seeds <- tf$unstack(tfp$random$split_seed(seed, n = 2L))
      n_leapfrog <- tf$random$stateless_uniform(
        shape = list(),
        seed = seeds[[1]],
        minval = self$l_min,
        maxval = self$l_max + 1L,
        dtype = tf$int32
      )
      accepted <- previous_kernel_results$accepted_results$`_replace`(
        num_leapfrog_steps = n_leapfrog
      )
      results <- previous_kernel_results$`_replace`(accepted_results = accepted)
      self$inner_kernel$one_step(current_state, results, seed = seeds[[2]])
    }
  )
)

# opt() with adam(), against tfp$math$minimize() running the same Keras
# optimiser with greta's convergence rule: stop once successive objectives
# differ by at most the tolerance. That rule is a ConvergenceCriterion
# subclass written in R. opt() starts from small random values on the free
# scale, as it does by default, and minimize() from zero. Both build and trace
# their loop afresh on every call, as opt() does.
AbsoluteChange <- reticulate::PyClass(
  "AbsoluteChange",
  inherit = tfp$optimizer$convergence_criteria$ConvergenceCriterion,
  defs = list(
    `__init__` = function(self, tolerance) {
      tfp$optimizer$convergence_criteria$ConvergenceCriterion$`__init__`(
        self,
        min_num_steps = 1L,
        name = "AbsoluteChange"
      )
      self$tolerance <- tolerance
      NULL
    },
    `_bootstrap` = function(self, loss, grads, parameters) {
      loss
    },
    `_one_step` = function(self, step, loss, grads, parameters, auxiliary_state) {
      reticulate::tuple(
        tf$less_equal(tf$abs(loss - auxiliary_state), self$tolerance),
        loss
      )
    }
  )
)

opt_tfp <- function(m, max_iterations, tolerance) {
  dag <- m$dag
  float <- greta:::tf_float()
  n_free <- length(greta:::unlist_tf(dag$example_parameters(free = TRUE)))
  free_state <- tf$Variable(matrix(0, 1, n_free), dtype = float)
  optimiser <- tf$keras$optimizers$Adam(
    learning_rate = bench_settings$opt_learning_rate
  )
  optimiser$build(list(free_state))
  # minimize() calls the loss with a seed, and falls back to calling it
  # without one only on Python's TypeError, which an R error is not
  loss <- function(seed = NULL) {
    tf$reshape(-dag$one_row_log_prob(free_state)$adjusted, shape = list())
  }
  run <- tensorflow::tf_function(function() {
    tfp$math$minimize(
      loss_fn = loss,
      num_steps = max_iterations,
      optimizer = optimiser,
      convergence_criterion = AbsoluteChange(tolerance = tolerance),
      trainable_variables = list(free_state)
    )
  })
  run()
  free_state$numpy()
}

bench_opt <- function(m) {
  bench::press(
    max_iterations = bench_settings$opt_max_iterations,
    {
      set.seed(session_seed)
      mark(
        greta = opt(
          m,
          optimiser = adam(learning_rate = bench_settings$opt_learning_rate),
          max_iterations = max_iterations,
          tolerance = 0
        ),
        tfp = opt_tfp(m, max_iterations, tolerance = 0),
        iterations = bench_settings$opt_calls,
        check = FALSE,
        memory = FALSE,
        filter_gc = FALSE
      )
    }
  )
}

The models

Two of greta’s example models: model-linear, a linear regression of attitude’s ratings on complaints, and model-hierarchical-linear, which regresses iris’s sepal length on sepal width with an offset for each species.

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)
}

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)
}

Session details

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

greta_sitrep()

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()

# 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)
 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)
 future        1.76.0     2026-09-24 [1] CRAN (R 4.6.1)
 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-07 [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)
 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)
 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)
 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/Rtmpa28bpm/file20404a580556
 [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

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.