Part VI. Software workflows in R, Stan, and JAX

Part VI

This part translates the theory and computation into concrete workflows in R, Stan, and JAX, with emphasis on reproducibility, diagnostics, and implementation discipline.

Part overview

Part VI translates the earlier theory into real computational workflow. The emphasis is not on software for its own sake, but on the disciplined use of tools to make Bayesian reasoning reproducible and diagnostically sound.

Part learning goals

By the end of Part VI, learners should be able to:

  • implement Bayesian analyses end-to-end in R;
  • write and diagnose Stan programs for nontrivial posterior inference;
  • use JAX for autodiff-based objectives, gradients, Hessians, and custom gradient-based Bayesian algorithms.

Fully written chapters in this repository

Reading path

Recommended order:

  1. Bayesian workflow in R
  2. Stan workflow for hierarchical models and HMC diagnostics
  3. JAX for gradient-based Bayesian computation

Chapter 24. Bayesian workflow in R

Core question

How can R serve as the central orchestration language for Bayesian simulation, summarization, visualization, and model criticism?

Section blueprint

  • Base R implementations of conjugate analysis.
  • Custom Metropolis and Gibbs samplers.
  • Posterior summaries and intervals.
  • Posterior predictive simulation.
  • Diagnostics using posterior, bayesplot, and loo.
  • Interfacing with cmdstanr.

Derivations and coding targets

  • Conjugate posterior functions as reusable R functions.
  • Hand-coded random-walk Metropolis for a simple target.
  • ESS and Monte Carlo standard error calculations from draws.
  • Plotting workflow for prior and posterior predictive checks.

Computational lab

Common pitfalls

  • Mixing data manipulation, model code, and diagnostics in one opaque script.
  • Reporting posterior summaries without retaining raw draws.
  • Forgetting reproducible seeds or session information.

Exercises

  • Write an R function for Gamma-Poisson updating.
  • Compute posterior predictive intervals from posterior draws.
  • Reproduce a bayesplot-style PPC using only base R.

Core references

  • R Core Team (2025)
  • Bürkner et al. (2024)
  • Gabry and Mahr (2024)
  • Vehtari et al. (2024)

Chapter 25. Bayesian workflow in Stan

Core question

How should a serious Bayesian analyst write, fit, diagnose, and revise Stan models?

Section blueprint

  • Stan language structure: data, parameters, transformed parameters, model, generated quantities.
  • Priors and parameter constraints.
  • HMC/NUTS computation through Stan.
  • Posterior predictive simulation.
  • Hierarchical models and reparameterization.
  • Diagnostics: divergences, \(\hat{R}\), ESS, E-BFMI, treedepth.
  • Practical tuning: adapt_delta, max_treedepth, initialization.

Derivations and coding targets

  • Translate a generative model into Stan blocks cleanly.
  • Compare centered and non-centered hierarchical parameterizations.
  • Use generated quantities for predictive simulation and pointwise log likelihood.
  • Read Stan diagnostics as geometry signals.

Computational lab

Common pitfalls

  • Placing deterministic transformations in the wrong Stan block.
  • Using weak priors that create poor geometry.
  • Treating tuning controls as substitutes for model redesign.

Exercises

  • Write a Stan program for Bayesian logistic regression with posterior predictive quantities.
  • Convert a centered hierarchical normal model into a non-centered form.
  • Explain what each Stan block contributes to the computation.

Core references

  • Carpenter et al. (2017)
  • Hoffman and Gelman (2014)
  • Betancourt (2017)

Full sample chapter

Chapter 26. Bayesian workflow in JAX

Core question

How can automatic differentiation and vectorized array programming be used to build custom Bayesian computation transparently?

Section blueprint

  • Why JAX belongs in a Bayesian course.
  • Autodiff, grad, hessian, vmap, and jit.
  • Writing log posterior functions.
  • MAP optimization and Laplace approximation.
  • Basic HMC-style mechanics in pure JAX.
  • Interoperability with broader probabilistic ecosystems.

Derivations and coding targets

  • Log posterior and gradient for Bayesian logistic regression.
  • Hessian-based curvature analysis around the MAP.
  • Basic leapfrog implementation.
  • Clear separation of model, gradient, optimizer, and sampler code.

Computational lab

Common pitfalls

  • Using JAX without understanding which functions are pure and differentiable.
  • Treating optimization convergence as posterior convergence.
  • Writing numerically unstable log-likelihood code.

Exercises

  • Implement a log posterior in JAX and verify gradients numerically.
  • Compare gradient ascent and Newton updates for the same target.
  • Use the Hessian at the MAP to build a Laplace approximation.

Core references

Full sample chapter

Back to top

References

Betancourt, Michael. 2017. “A Conceptual Introduction to Hamiltonian Monte Carlo.” arXiv Preprint arXiv:1701.02434.
Bradbury, James, Roy Frostig, Peter Hawkins, et al. 2018. “JAX: Composable Transformations of Python+NumPy Programs.” GitHub Repository.
Bürkner, Paul-Christian, Jonah Gabry, Aki Vehtari, et al. 2024. Posterior: Tools for Working with Posterior Distributions.
Carpenter, Bob, Andrew Gelman, Matthew D. Hoffman, et al. 2017. “Stan: A Probabilistic Programming Language.” Journal of Statistical Software 76 (1): 1–32.
Gabry, Jonah, and Tristan Mahr. 2024. Bayesplot: Plotting for Bayesian Models.
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.
Murphy, Kevin P. 2023. Probabilistic Machine Learning: Advanced Topics. MIT Press.
Phan, Du, Neeraj Pradhan, and Martin Jankowiak. 2020. “Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro.” arXiv Preprint arXiv:1912.11554.
R Core Team. 2025. R: A Language and Environment for Statistical Computing. R Foundation for Statistical Computing.
Vehtari, Aki, Jonah Gabry, Yuling Yao, and Andrew Gelman. 2024. Loo: Efficient Leave-One-Out Cross-Validation and WAIC for Bayesian Models.