# Automatically Batching Control-Intensive Programs for Modern Accelerators

arXiv: https://arxiv.org/abs/1910.11141 · HTML: https://mlsanity.com/paper/1910.11141/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 23 Oct 2019.

## Built on by (most cited first, 6 of 6)

- [Disentangling impact of capacity, objective, batchsize, estimators, and step-size on flow VI](https://mlsanity.com/paper/2412.08824/) (2024)
- [Incorporating the ChEES Criterion into Sequential Monte Carlo Samplers](https://mlsanity.com/paper/2504.02627/) (2025)
- [ACRoBat: Optimizing Auto-batching of Dynamic Deep Learning at Compile Time](https://mlsanity.com/paper/2305.10611/) (2023)
- [Efficiently Vectorized MCMC on Modern Accelerators](https://mlsanity.com/paper/2503.17405/) (2025)
- [Differentiate the Evaluator, Not the Program: An Efficient Runtime Representation for Neuro-Symbolic Learning](https://mlsanity.com/paper/2607.03574/) (2026)
- [DrainSinkhorn: Safe Elimination for Batched Entropic Optimal Transport](https://mlsanity.com/paper/2607.24741/) (2026)

## Abstract (arXiv)

We present a general approach to batching arbitrary computations for accelerators such as GPUs. We show orders-of-magnitude speedups using our method on the No U-Turn Sampler (NUTS), a workhorse algorithm in Bayesian statistics. The central challenge of batching NUTS and other Markov chain Monte Carlo algorithms is data-dependent control flow and recursion. We overcome this by mechanically transforming a single-example implementation into a form that explicitly tracks the current program point for each batch member, and only steps forward those in the same place. We present two different batching algorithms: a simpler, previously published one that inherits recursion from the host Python, and a more complex, novel one that implemenents recursion directly and can batch across it. We implement these batching methods as a general program transformation on Python source. Both the batching system and the NUTS implementation presented here are available as part of the popular TensorFlow Probability software package.