# Ideas for Testing Randomized Algorithms
> Source:
> Published: 2026-11-04 08:57:20+00:00
Randomized algorithms are hard to test. Their outputs are, by design,
unpredictable—and that unpredictability is a *feature*, not a bug. It is what
gives them useful guarantees like unbiasedness and computational efficiency. But
it also makes them slippery: how do you check that a stochastic program is
correct when it gives different answers every time you run it?
This post sketches some ideas for how to think about testing randomized algorithms—especially the kinds of algorithms that we use for probabilistic inference. I'll present a concrete sampling algorithm with several plausible implementations—some correct, some subtly broken—and develop progressively more powerful testing strategies to tell them apart. Each strategy catches a bug that the previous one missed.
The algorithm is a small building block. Its job is to propose the next token inside sequential Monte Carlo (SMC) for constrained generation from language models. Testing the block in isolation is only half the job, so at the end we'll plug it into SMC and test the whole pipeline.
## A Tricky Little Sampling Algorithm
Suppose we have a distribution $p$ over a finite set $\mathcal{X} = \{1,
\ldots, N\}$ and a condition $C \subseteq \mathcal{X}$ (the set of "good"
elements). We want to sample from the *conditional* distribution $p(\cdot \mid
C)$. But rather than returning a plain sample, we want a *properly weighted*
sample $(x, w)$—a sample-weight pair that can be used for unbiased estimation.
The following algorithm, which we call AWRS (from [Lipkin et al., 2025](https://arxiv.org/abs/2504.05410)), does exactly this. Here is the pseudocode from the paper:
The idea: run rejection sampling twice. The first pass finds a conditioned sample $x_c$ while tracking the total rejected probability mass $\psi_0$. The second pass generates additional rejections to get a count $n$. The weight $(1 - \psi_0) / (n + 1)$ combines these two signals—the rejected mass and the rejection count—into an unbiased estimate of the normalizing constant $Z$. The number of rejections before finding a valid token contains signal about the magnitude of $Z$, and weighting by the actual probability mass removed (rather than just counting) reduces variance.
### Four Implementations, at Least One Bug
Did we implement it correctly? That's the question we'd like our tests to answer. To make the exercise concrete, here are four Python implementations. They all look reasonable, but at least one has a bug. Which ones are correct?
Each variant takes a `draw` parameter that defaults to sampling from a
categorical distribution. This is a hook we'll exploit later: in production it
calls the RNG, but in testing we'll swap it for something that lets us
enumerate every possible execution.
## Variant A
``` python
def awrs_a(p, cond, draw=sample_categorical):
p_cond = p.copy()
n, psi = 0, 0
xc0 = None
for i in range(2):
while True:
q = p_cond / p_cond.sum()
x = draw(q, context=('awrs', i, xc0))
if cond(x):
if i == 0:
xc0 = x
break
else:
if i == 0:
psi += p_cond[x]
p_cond[x] = 0
n += 1
return xc0, (1 - psi) / (n + 1)
```
## Variant B
``` python
def awrs_b(p, cond, draw=sample_categorical):
p_cond = p.copy()
n, psi = 0, 0
xc0 = None
for i in range(2):
while True:
q = p_cond / p_cond.sum()
x = draw(q, context=('awrs', i, xc0))
if cond(x):
if i == 0:
xc0 = x
break
else:
if i == 0:
psi += p_cond[x]
p_cond[x] = 0
n += 1
return xc0, (1 + psi) / (n + 1)
```
## Variant C
``` python
def awrs_c(p, cond, draw=sample_categorical):
p_cond = p.copy()
n, psi = 0, 0
xc0 = None
for i in range(2):
while True:
q = p_cond / p_cond.sum()
x = draw(q, context=('awrs', i, xc0))
if cond(x):
if i == 0:
xc0 = x
break
else:
psi += p_cond[x]
p_cond[x] = 0
n += 1
return xc0, (1 - psi) / (n + 1)
```
## Variant D
``` python
def awrs_d(p, cond, draw=sample_categorical):
p_cond = p.copy()
n, psi = 0, 0
xc0 = None
for i in range(2):
while True:
q = p_cond / p_cond.sum()
x = draw(q, context=('awrs', i, xc0))
if cond(x):
if i == 0:
xc0 = x
break
else:
if i == 0:
psi += q[x]
p_cond[x] = 0
n += 1
return xc0, (1 - psi) / (n + 1)
```
The variants differ from one another by a line or two. I encourage you to read them carefully and try to spot the differences before scrolling further.
## Hint
Compare the `else` branch and the `return` statement across the four variants.
Our job is to build testing strategies that catch the buggy ones—ideally without needing to know which line is wrong.
But before we can test *anything*, we need to answer the most important
question in all of software testing:
### What Should This Code Do?
You cannot test code without a specification. For randomized algorithms, the specification comes from theory—and translating theory into a checkable property is half the battle. We'll keep coming back to this question for every algorithm in this post.
For AWRS, the specification is: it should return a sample-weight pair $(x, w)$
such that $w \cdot \mathbf{1}[x = x']$ is an **unbiased estimator** of the
unnormalized conditional distribution for every $x' \in \mathcal{X}$:
$$ \mathop{\mathbb{E}}_{(x,w) \sim \text{AWRS}} \left[ w \cdot \mathbf{1}[x = x'] \right] = p(x') \cdot \mathbf{1}[x' \in C] $$
The right-hand side is the unnormalized target—$p$ restricted to the condition $C$, with normalizing constant $Z = \sum_{x \in C} p(x)$.
That AWRS satisfies this is a theorem in
[Lipkin et al., 2025](https://arxiv.org/abs/2504.05410), and we'll take it as
given. Our question is whether the code does what the paper proves.
Note that you may often need to tweak definitions, narrow or broaden the scope of theory in order to suit your code's intended use cases. This is an important skill for turning theory into tests.
### Setting Up a Test Case
Now that we have a testable definition, we need to think about which things in the equation we can compute and which things we need to estimate. Let's make a concrete test case:
```
p = np.array([0.5, 0.3, 0.15, 0.04, 0.01]) # a distribution with a long tail
cond = np.array([True, False, True, False, False]) # the "good" elements
N = len(p)
Z = p[cond].sum()
target = p * cond # unnormalized target (not a distribution!)
```
## Strategy 1: Monte Carlo Averaging
The obvious first testing strategy: run the code "a lot of times" and compute an average.
```
np.random.seed(0)
reps = 1_000
for name, awrs in variants.items():
have = np.zeros(N)
for r in range(1, 1 + reps):
(x, w) = awrs(p, lambda x: cond[x])
have[x] += w
err = abs(have / reps - target).sum()
print(f'Variant {name}: error = {err:.6f}')
Variant A: error = 0.023367
Variant B: error = 0.098657
Variant C: error = 0.030520
Variant D: error = 0.012054
```
Is this error small enough? How do you know? (Variant D has the smallest error of the four.)
This is the central problem with Monte Carlo testing: a single number tells you
almost nothing without a reference frame. We need to know what *should* happen
as we increase the number of samples.
### Convergence Rates
A better approach is to watch the error decrease as a function of the number of
samples $M$. It should keep going down *forever*, modulo numerical precision.
And crucially, we often know the *rate* at which it should decrease.
The general picture: suppose we have a sequence of estimators $\widehat{\mu}^{(M)}$ targeting some quantity $\mu$. The mean squared error decomposes as
$$ \mathop{\mathbb{E}} \| \widehat{\mu}^{(M)} - \mu \|^2 = \| \mathop{\mathbb{E}}[\widehat{\mu}^{(M)}] - \mu \|^2 + \mathop{\mathbb{E}} \| \widehat{\mu}^{(M)} - \mathop{\mathbb{E}}[\widehat{\mu}^{(M)}] \|^2 $$
The first term is the squared bias; the second is the variance (more precisely, the trace of the covariance of $\widehat{\mu}^{(M)}$). Different classes of estimators make different predictions about how these terms behave:
-
**I.I.D. unbiased** (e.g., our AWRS test). The bias is zero for all $M$.
The variance is $V/M$ where $V = \text{tr}(\text{Cov}(X))$—the sum of
the componentwise variances of a single sample. So the mean squared error is
exactly $V/M$.If you prefer root mean squared
error, the slope is $-1/2$ and the intercept is $\frac{1}{2} \log
V$.
-
**Correlated but unbiased** (e.g., sampling without replacement). The bias
is still zero. Samples drawn without replacement are negatively correlated,
which shrinks the variance: the rate is still $O(1/M)$, but the constant is
smaller.
-
**Biased but consistent** (e.g., self-normalized importance sampling, MCMC). The
bias is $O(1/M)$, so the squared bias is $O(1/M^2)$. The variance is
$O(1/M)$.
The MSE is still $O(1/M)$ overall, so on a log–log plot the slope looks
the same—the bias hides in the noise at large $M$.
In all three cases, $\text{MSE} \sim c / M$ for some constant $c$, which means $\log \text{MSE} = -\log M + \log c$: a line with slope $-1$ on a log–log plot. The intercept $\log c$ tells us the constant.
## Deriving $\,V\,$ for vector-valued estimators
Our estimator is vector-valued (a distribution over $N$ elements), so we need a scalar measure of error. The natural choice is the squared $L_2$ norm $\| \widehat{\mu}^{(M)} - \mu \|^2$ where $\widehat{\mu}^{(M)} = \frac{1}{M} \sum_{m=1}^{M} X^{(m)}$ and $\mu = \mathop{\mathbb{E}}[X^{(m)}]$:
$$ \mathop{\mathbb{E}} \| \widehat{\mu}^{(M)} - \mu \|^2 = \mathop{\mathbb{E}} \sum_{j=1}^{N} (\widehat{\mu}^{(M)}_j - \mu_j)^2 = \sum_{j=1}^{N} \text{Var}(\widehat{\mu}^{(M)}_j) = \frac{1}{M} \sum_{j=1}^{N} \text{Var}(X^{(m)}_j) = \frac{V}{M} $$
where $V = \text{tr}(\text{Cov}(X^{(m)}))$ is the trace of the covariance matrix of a single sample. The last step uses independence; for correlated samples, the $1/M$ factor becomes $1/M$ times an autocorrelation correction.
```
results = {}
for name, awrs in variants.items():
np.random.seed(0)
have = np.zeros(N)
xs, ys = [], []
for r in range(1, 1 + 1_000):
(item, w) = awrs(p, lambda x: cond[x])
have[item] += w
if r >= 10:
xs.append(r)
ys.append(abs(have / r - target).sum())
slope, intercept = np.polyfit(np.log(xs), np.log(ys), deg=1)
results[name] = (xs, ys, slope, intercept)
plot_single_runs(results)
```
A note on reading this plot: when we say "log–log scale" we mean the axes are
log-transformed but the tick labels show the original values (e.g., $10^2$,
$10^3$). This is what a log axis does in a plotting library. This is different from
*literally plotting* $\log y$ vs $\log x$, where the axes are linear but the
values are logs. The line-fitting code does the latter (`np.polyfit` on
$\log x$, $\log y$), but for visual inspection, log-scaled axes are usually
clearer because you can read off the original quantities directly.
For a correct estimator, the fitted slope should be close to $-0.5$ (we're
plotting error, not squared error). Variant B fails this test plainly: its
error stops going down after a few dozen samples and the slope is positive.
**B is out.**
The other three look fine. C's slope is a little shallow, but the slope of a single run is itself a noisy statistic—rerunning the correct variant with other seeds gives slopes anywhere from about $-0.25$ to $-0.75$. D's curve is indistinguishable from A's. That's not a coincidence, and we'll see why later.
### What a Single Run Can and Can't Tell Us
Watching a single long run tells us the *rate* of convergence, but it can't
distinguish the three cases above. An unbiased i.i.d. estimator, an estimator
built from samples drawn without replacement, and an MCMC chain can all produce
very similar log–log plots—the slopes are the same and the constants differ in ways
that are hard to eyeball.
A pedagogical example: suppose we accidentally computed the average as
$\frac{1}{M+1} \sum_{m=1}^{M} X^{(m)}$ instead of $\frac{1}{M} \sum_{m=1}^{M}
X^{(m)}$. This estimator is *biased*—it systematically underestimates the
target—but *consistent*, because the bias factor $M/(M+1) \to 1$. The
single-run convergence plot would look nearly identical.
The same thing is happening to C. Its bias is real, but at a thousand samples it is still smaller than the noise. A longer run would expose it eventually, but how long is long enough? A single run has no way to tell us.
### Averaging Over Multiple Runs
To separate bias from variance, we need **multiple independent runs**. The idea:
run the estimator many times at each sample size $M$, and compare the *average*
behavior to the theoretical prediction.
Recall the decomposition from earlier:
$$ \mathop{\mathbb{E}} \| \widehat{\mu}^{(M)} - \mu \|^2 = \| \mathop{\mathbb{E}}[\widehat{\mu}^{(M)}] - \mu \|^2 + \mathop{\mathbb{E}} \| \widehat{\mu}^{(M)} - \mathop{\mathbb{E}}[\widehat{\mu}^{(M)}] \|^2 $$
For the i.i.d. unbiased case, the variance term is $V/M$. So:
- If the observed MSE matches $V/M$, the squared bias must be near zero—evidence of unbiasedness.
- If the estimator has persistent bias, $\| \mathop{\mathbb{E}}[\widehat{\mu}^{(M)}] -
\mu \|^2$ does not vanish with $M$. At large $M$, the variance $V/M$ shrinks
but the bias dominates, causing the observed curve to *flatten out* above the
theoretical line.
- For biased-but-consistent estimators (like self-normalized IS), the squared
bias is $O(1/M^2)$—it vanishes
faster than the variance. This means the MSE curve will *eventually* match
$V/M$, but may sit slightly above it for small $M$.
```
results = {}
for name, awrs in variants.items():
np.random.seed(0)
reps, M = 1000, 200
errors = np.zeros((reps, M)) # squared error of the running average, per run
samples = [] # single-sample estimates, for V
for r in range(reps):
have = np.zeros(N)
for m in range(1, M + 1):
(item, w) = awrs(p, lambda x: cond[x])
have[item] += w
onehot = np.zeros(N)
onehot[item] = w
samples.append(onehot)
errors[r, m - 1] = np.linalg.norm(have / m - target)**2
V = np.trace(np.cov(np.array(samples).T))
results[name] = (errors, V)
plot_mse(results)
```
A and D hug the theoretical $V/M$ line, which is strong evidence that they are
properly weighted—not just consistent, but unbiased, with the right
variance. B flattens out, as before. C is the interesting one: it tracks the
line for small $M$ and then lifts off as the variance shrinks below its squared
bias. Averaging over a thousand runs makes the gap far larger than the
confidence band. **C is out.**
D, on the other hand, has now passed every test we've thrown at it.
But Monte Carlo testing is still a *blunt instrument*. Small biases hide
easily in the noise, and you need many runs to have any statistical power. For
correlated estimators, the effective variance may be hard to compute, making the theoretical line itself uncertain. In software—unlike
the real world—we can almost always do better.
## Strategy 2: Program Tracing
Here's the key insight: if the state space is small enough, we can *enumerate
all possible executions* of the randomized program and verify the correctness
condition *exactly* (up to floating point).
The idea is to hijack the random number generation inside the program to make it systematically explore every possible branch. We track each execution trace's probability, and combine all traces to compute the exact output distribution.
### The Tracer
The tracing strategy comes from
[Shi et al., 2020](https://arxiv.org/abs/2002.09067), who use it to sample
without replacement from programs such as sequence models. The same idea
underlies exact inference by enumeration in probabilistic programming
languages: WebPPL explores every execution of a program by capturing a
continuation at each random choice
([Goodman & Stuhlmüller, 2014](http://dippl.org/chapters/03-enumeration.html)).
The implementation uses a trie data structure that mirrors the branching structure of the
program's random choices. Each node in the trie corresponds to a decision
point, and its children correspond to the possible outcomes.
Each node stores the probability mass that hasn't been explored below it. A
node is created the first time the program reaches it, holding the probability
of the path that leads there, so at the end of a trace `tracer.cur.mass` is
that trace's probability; this is how the checks below weight each output. When
a trace finishes, its mass is subtracted from every node on the way back up to
the root. Choosing children in proportion to their unexplored mass then never
revisits a finished trace, and the root's mass is the total probability of the
traces not yet run.
Here is the tracer at work on variant A with our test case. Each step runs the
program once more and adds one trace to the trie. Node colors mark which pass of
AWRS made the choice, and a node fades as the probability mass below it is
explored. The order menu is explained [below](#which-trace-next).
The beautiful thing about `TraceSWOR` is that it works as a context manager:
each `with tracer:` block explores one trace, and when it exits, that trace's
mass is zeroed out so we never revisit it. We loop until the root's mass hits
zero—meaning we've exhausted every possible execution.
The tracer is also available as a small, tested library,
[`program_tracing`](https://github.com/timvieira/testing-randomized-algorithms),
in this post's repository. It includes a `traces(program)` helper that runs this
loop for you.
## Tracer implementation
``` python
class Node:
"""A node in the execution trie."""
def __init__(self, idx, mass, parent, token=None):
self.idx = idx
self.mass = mass # remaining probability mass
self.parent = parent
self.token = token # label for visualization
self.child_masses = None # set on first visit
self.keys = None # outcomes of the choice made here
self.active_children = {}
self._mass = mass # remember original mass
self.context = None
def sample(self):
"""Sample a child proportional to remaining mass."""
p = self.child_masses / self.child_masses.sum()
return np.random.choice(len(p), p=p)
def update(self):
"""Propagate mass changes up to the root."""
if self.parent is not None:
self.parent.child_masses[self.idx] = self.mass
self.parent.mass = np.sum(self.parent.child_masses)
self.parent.update()
class Tracer:
"""
Lazily materializes the probability tree of a generative process
by program tracing. Works with both array-valued and dict-valued
distributions.
"""
def __init__(self):
self.root = Node(idx=-1, mass=1.0, parent=None)
self.cur = None
def __call__(self, p, context=None):
"""Drop-in replacement for any sampling function."""
if isinstance(p, dict):
keys = list(p.keys())
vals = np.array([p[k] for k in keys], dtype=float)
else:
p = np.asarray(p, dtype=float)
keys = list(range(len(p)))
vals = p
vals = vals / vals.sum()
cur = self.cur
if cur.child_masses is None:
cur.child_masses = cur.mass * vals
cur.keys = keys
cur.context = context
elif keys != cur.keys:
raise ValueError('the program made a different choice at the same point '
'of two traces; it must be deterministic given its draws')
a = self.choose(cur)
if a not in cur.active_children:
cur.active_children[a] = Node(
idx=a, mass=cur.child_masses[a], parent=cur,
token=keys[a],
)
self.cur = cur.active_children[a]
return keys[a]
def choose(self, node):
"""Pick which child to explore next."""
return node.sample()
class TraceSWOR(Tracer):
"""
Sampling without replacement meets program tracing.
Each execution trace is explored exactly once. After a trace
completes, its probability mass is zeroed out and propagated
up the trie, ensuring the next trace is fresh.
"""
def __enter__(self):
self.cur = self.root
return self
def __exit__(self, *args):
self.cur.mass = 0
self.cur.update()
```
### Exact Verification
Now we can verify all four AWRS variants *exactly*:
``` python
def exact_moments(awrs):
"""Exact first and second moments of w * onehot(x), over all traces."""
tracer = TraceSWOR()
m1, m2 = np.zeros(N), np.zeros(N)
while tracer.root.mass > 0:
with tracer:
(item, w) = awrs(p, lambda x: cond[x], draw=tracer)
m1[item] += w * tracer.cur.mass
m2[item] += w**2 * tracer.cur.mass
return m1, m2
for name, awrs in variants.items():
prob, _ = exact_moments(awrs)
err = abs(target - prob).sum()
status = '\u2713' if np.allclose(target, prob) else '\u2717'
print(f'Variant {name}: {status} (error = {err:.2e})')
Variant A: ✓ (error = 3.61e-16)
Variant B: ✗ (error = 9.08e-02)
Variant C: ✗ (error = 3.15e-02)
Variant D: ✗ (error = 1.61e-04)
```
This is not a statistical test. It is an *exact* check (modulo floating point)
that the proper weighting condition holds.The difference between `target` and `prob` is not zero—it rarely will be with floating point numbers. But `np.allclose` checks that they agree to within a reasonable tolerance ($\approx 10^{-8}$). No confidence intervals, no
squinting at log–log plots. The tracer gives a definitive yes or no for each
variant. **D is out**, and A is the only survivor.
D's bug is in the `else` branch: it accumulates the *normalized* probability
`q[x]` instead of the remaining mass `p_cond[x]`. The two agree on the first
rejection, because `p_cond` still sums to one. They differ only when the first
pass rejects two or more elements. With our long-tailed `p`, that happens in
about one run in twenty, and even then the weight is off by only a couple of
percent[✓](https://github.com/timvieira/testing-randomized-algorithms/blob/master/tests/test_article.py#L85)—which is why D's single run looked just like A's.
### How Many Samples Would Monte Carlo Need?
The tracer computes more than a pass/fail bit. Collecting the second moment alongside the first gives the exact squared bias and the exact variance $V$ of each variant. The bias of the average stays fixed as $M$ grows while its variance shrinks like $V/M$, so the two cross at $M^* = V / \text{bias}^2$. Below $M^*$, the bug is buried in the noise.
```
for name, awrs in variants.items():
m1, m2 = exact_moments(awrs)
bias2 = np.sum((m1 - target)**2)
V = np.sum(m2 - m1**2)
if np.isclose(bias2, 0, atol=1e-20):
print(f'Variant {name}: unbiased')
else:
print(f'Variant {name}: bias^2 = {bias2:.1e}, V = {V:.2f}, '
f'M* = {V / bias2:,.0f}')
Variant A: unbiased
Variant B: bias^2 = 5.3e-03, V = 0.25, M* = 47
Variant C: bias^2 = 6.4e-04, V = 0.24, M* = 382
Variant D: bias^2 = 1.7e-08, V = 0.24, M* = 14,544,122
```
This is the ladder we just climbed. B surfaces after a few dozen samples, so a single run catches it. C needs a few hundred, so it takes careful averaging. D needs about fifteen million, and no reasonable Monte Carlo test would ever see it. The tracer took a fraction of a second.
### Comparing Variances
The same moments let us compare estimators that are all correct. Unbiasedness is only half the story: AWRS exists because it has lower variance than simpler alternatives. The simplest properly weighted sampler draws once from $p$ and keeps the draw with weight 1 if it lands in $C$, or weight 0 otherwise. It is unbiased, so no correctness test can tell it apart from AWRS. The tracer can compare their variances exactly:
``` python
def naive(p, cond, draw=sample_categorical):
"""Draw once from p; weight 1 if the draw satisfies the condition, else 0."""
x = draw(p)
return x, float(cond(x))
for name, sampler in [('AWRS', awrs_a), ('naive', naive)]:
m1, m2 = exact_moments(sampler)
bias2 = np.sum((m1 - target)**2)
V = np.sum(m2 - m1**2)
var_w = m2.sum() - m1.sum()**2 # variance of w as an estimate of Z
print(f'{name:>5}: bias^2 = {bias2:.1e}, V = {V:.3f}, Var(w) = {var_w:.3f}')
AWRS: bias^2 = 6.5e-32, V = 0.244, Var(w) = 0.094
naive: bias^2 = 0.0e+00, V = 0.378, Var(w) = 0.227
```
Both are unbiased, but the naive weight has about 2.4 times the variance of the
AWRS weight as an estimate of $Z$, so it needs about 2.4 times as many samples
for the same accuracy.[✓](https://github.com/timvieira/testing-randomized-algorithms/blob/master/tests/test_article.py#L121) This turns a performance claim into a test. If a
refactor leaves an estimator correct but noisier, an exact variance check
catches it, with no line to fit through a noisy plot. The same goes for the
SMC pipeline below: the tracer gives the exact variance of $\widehat{Z}$, so we
can check whether a better proposal or more particles actually help.
### Which Trace Next?
`TraceSWOR` makes each choice at random, in proportion to the unexplored mass
below it. That isn't the only reasonable order. A deterministic alternative is
**greedy**: always descend into the child with the most unexplored mass. It's a
one-method subclass:
```
class TraceGreedy(TraceSWOR):
"""Always descend into the child with the most unexplored mass."""
def choose(self, node):
return int(np.argmax(node.child_masses))
```
If we run every trace, the order doesn't matter: both give the same exact
answer. It matters when we stop early. After exploring traces with total
probability $c$, the unexplored traces have total probability $1 - c$, so if
every weight is at most $B$, the answer can be off by at most $B(1 - c)$.
(AWRS's weights are at most 1.[✓](https://github.com/timvieira/testing-randomized-algorithms/blob/master/tests/test_article.py#L131)) The bound holds for any order, and greedy
tightens it faster because it heads for the heaviest unexplored part of the
trie:
```
program = lambda draw: awrs_a(p, lambda x: cond[x], draw=draw)
def trace_masses(tracer):
masses = []
while tracer.root.mass > 0:
with tracer:
program(tracer)
masses.append(tracer.cur.mass)
return masses
np.random.seed(0)
orders = {
'sampling': trace_masses(TraceSWOR()),
'greedy': trace_masses(TraceGreedy()),
}
plot_unexplored(orders)
```
Sampling covers mass more slowly, but it offers something greedy doesn't. The
traces it has seen are a sample without replacement, which supports an unbiased
estimate of the part it hasn't seen
([Shi et al., 2020](https://arxiv.org/abs/2002.09067)). A deterministic order
gives a guaranteed bound on the error; sampling gives an unbiased estimate. The
replay above can run in either order.
Why not go further, and always run the most probable unexplored trace next?
Greedy heads for the heaviest subtree, not the heaviest single trace, so it
isn't that. True best-first search
needs to resume the program at an arbitrary unexplored branch, and the tracer
can't: it remembers probabilities, not the program's state at each choice
point. The workaround is to rerun the program from the start, forcing the
choices along the path to that branch, and stop it there. That adds
assumptions on top of the ones tracing already makes—rerunning must have
no side effects, and the program must be safe to interrupt partway—and it
costs a rerun for every node of the trie rather than one per trace.
Probabilistic programming languages avoid the rerun by capturing a
continuation at each choice point. WebPPL keeps them in a priority queue, which
gives exactly this most-likely-first order
([Goodman & Stuhlmüller, 2014](http://dippl.org/chapters/03-enumeration.html)).
Without continuations, it's not something to bolt onto a tracer casually.
## Strategy 3: Testing the Whole Pipeline
AWRS is not meant to be used on its own. In [Lipkin et al., 2025](https://arxiv.org/abs/2504.05410),
it is the proposal step inside sequential Monte Carlo for generating text from
a language model subject to a constraint. Each particle asks AWRS for a next
token that satisfies the constraint, and multiplies its weight by the weight
that AWRS returns.
Program tracing is not limited to a single call of a sampling routine. It works on anything that makes discrete random choices through a hookable interface, so we can trace the whole pipeline at once—every AWRS call in every particle, plus the resampling steps.
### What Should This Code Do?
Let $p$ be a language model over strings, and let $C$ be a constraint that can be checked one token at a time. The unnormalized target is
$$ \gamma(\mathbf{x}) = p(\mathbf{x}) \cdot \mathbf{1}[\mathbf{x} \in C] $$
and its normalizing constant $Z = \sum_{\mathbf{x}} \gamma(\mathbf{x})$ is the probability that a string drawn from $p$ satisfies the constraint. SMC returns $M$ weighted particles $(\mathbf{x}_m, w_m)$, and the specification is the proper weighting condition from before, lifted from tokens to strings:
$$ \mathop{\mathbb{E}} \left[ \frac{1}{M} \sum_{m=1}^{M} w_m \cdot \mathbf{1}[\mathbf{x}_m = \mathbf{x}'] \right] = \gamma(\mathbf{x}') $$
for every string $\mathbf{x}'$. Summing over $\mathbf{x}'$ gives the more familiar statement that $\frac{1}{M} \sum_m w_m$ is an unbiased estimate of $Z$. Checking every string is the stronger test: a bug that moves mass from one string to another without changing the total would slip past a check on $Z$ alone.
Why should this hold? Proper weighting composes. At step $t$, AWRS returns a token and a weight that are properly weighted for $p(\cdot \mid \mathbf{x}_{\lt t})$ restricted to the allowed tokens. Multiplying these weights along a string gives a properly weighted sample for the product of the per-step targets, which is $\gamma$. Resampling preserves proper weighting as long as every resampled particle carries the average weight.
### Setting Up a Test Case
The language model is a small, explicit distribution over strings. The
constraint says that a `b` must be followed by an `a` or by the end of the
string, so `bb` and `bc` are ruled out:
```
lm = SimpleLM({
(): 0.1, ('a',): 0.1, ('a', 'a'): 0.1, ('a', 'b'): 0.1,
('b',): 0.1, ('b', 'a'): 0.15, ('b', 'b'): 0.25, ('b', 'c'): 0.1,
})
VOCAB = ['a', 'b', 'c', EOS]
def allowed(prefix, token):
"""After a 'b', the next token must be 'a' or EOS."""
return not (prefix[-1:] == ('b',) and token not in ('a', EOS))
target_lm = {s: prob * is_valid(s) for s, prob in lm.sequences.items()}
Z_lm = sum(target_lm.values()) # = 0.65, not 1
```
## Utilities: `SimpleLM`, `is_valid`
```
EOS = '' # end-of-sequence marker
class SimpleLM:
"""
A sequence model defined by an explicit, finite distribution over strings.
Maps prefixes to next-token distributions (as dicts).
"""
def __init__(self, sequences):
"""
sequences: dict mapping tuple-of-tokens to probability.
E.g., {('a','a'): 0.3, ('b','b','b'): 0.6, (): 0.1}
"""
self.sequences = sequences
def p_next(self, prefix):
"""Return conditional distribution over next token given prefix."""
dist = {}
for seq, prob in self.sequences.items():
if len(seq) < len(prefix):
continue
if seq[:len(prefix)] != prefix:
continue
if len(seq) == len(prefix):
dist[EOS] = dist.get(EOS, 0) + prob
else:
token = seq[len(prefix)]
dist[token] = dist.get(token, 0) + prob
total = sum(dist.values())
return {k: v / total for k, v in dist.items()}
def is_valid(string):
return all(allowed(string[:t], string[t]) for t in range(len(string)))
```
Here is the sampler. At each step, every unfinished particle extends itself
with AWRS and multiplies in the AWRS weight. When the effective sample size
drops below `tau` times the number of live particles, it resamples. With
`tau=0` it never resamples, and the algorithm reduces to sequential importance
sampling (SIS); with `tau>0` it is sequential importance resampling (SIR).
``` python
from collections import namedtuple
Particle = namedtuple('Particle', 'weight, string')
def smc(lm, M, tau, awrs, draw=sample_categorical):
unfinished = [Particle(1, ()) for _ in range(M)]
finished = []
while unfinished:
next_unfinished = []
for particle in unfinished:
P = lm.p_next(particle.string)
p = np.array([P.get(v, 0) for v in VOCAB])
i, w = awrs(p, lambda i: allowed(particle.string, VOCAB[i]), draw=draw)
x = VOCAB[i]
if x == EOS:
finished.append(Particle(particle.weight * w, particle.string))
else:
next_unfinished.append(Particle(particle.weight * w,
particle.string + (x,)))
unfinished = next_unfinished
if not unfinished:
break
W = np.array([particle.weight for particle in unfinished])
if W.sum()**2 / (W** 2).sum() < tau * len(unfinished):
indices = [draw(W / W.sum(), context='resample') for _ in unfinished]
unfinished = [Particle(W.sum() / len(unfinished), unfinished[k].string)
for k in indices]
return finished
```
The same `draw` hook is passed down to every AWRS call and to the resampling
step, so a single tracer sees every random choice the pipeline makes.
### Sequential Importance Sampling
One particle is enough to check the proper weighting condition for SIS. We run
it with each AWRS variant, plus one more contender: *locally constrained
decoding*, which samples each token from the language model restricted to the
allowed tokens and ignores the weight. It is the obvious way to do constrained
generation.
``` python
def locally_constrained(p, cond, draw=sample_categorical):
"""AWRS's sample, but with its weight thrown away."""
x, _ = awrs_a(p, cond, draw)
return x, 1
def check_smc(awrs, M, tau):
tracer = TraceSWOR()
have = defaultdict(float)
n_traces = 0
while tracer.root.mass > 0:
with tracer:
for particle in smc(lm, M, tau, awrs, draw=tracer):
have[particle.string] += particle.weight / M * tracer.cur.mass
n_traces += 1
err = sum(abs(have[s] - target_lm[s]) for s in target_lm)
status = '\u2713' if err < 1e-10 else '\u2717'
return status, err, n_traces
print(f'Z = {Z_lm:.2f}')
contenders = {**variants, 'locally constrained': locally_constrained}
for name, awrs in contenders.items():
status, err, n_traces = check_smc(awrs, M=1, tau=0)
print(f'SIS with {name}: {status} '
f'(error = {err:.2e}, traces = {n_traces})')
Z = 0.65
SIS with A: ✓ (error = 2.78e-17, traces = 162)
SIS with B: ✗ (error = 1.21e-01, traces = 162)
SIS with C: ✗ (error = 3.34e-02, traces = 162)
SIS with D: ✗ (error = 4.22e-03, traces = 162)
SIS with locally constrained: ✗ (error = 3.50e-01, traces = 162)
```
Every bug is caught through the pipeline, including D. After a `b`, two of
the four tokens are forbidden, so AWRS's first pass can reject twice in one
step—exactly the situation that triggers D's bug.
The last line is the most instructive one. Locally constrained decoding never
produces an invalid string, but it gets the probabilities wrong: it
over-represents strings whose prefixes leave little room under the constraint.
Here, it puts too much mass on `b` and `ba`[✓](https://github.com/timvieira/testing-randomized-algorithms/blob/master/tests/test_article.py#L163), because it never pays for the
probability that `b` would have been followed by a forbidden token. This bias
is the reason AWRS returns a weight in the first place.
### Sequential Importance Resampling
Resampling is just one more random choice, made through the same `draw` hook,
so the tracer needs no changes. Two particles are enough to exercise it:
```
for name in ['A', 'D']:
status, err, n_traces = check_smc(variants[name], M=2, tau=0.75)
print(f'SIR with {name}: {status} '
f'(error = {err:.2e}, traces = {n_traces:,})')
SIR with A: ✓ (error = 1.10e-14, traces = 28,836)
SIR with D: ✗ (error = 4.22e-03, traces = 30,132)
```
The correct implementation passes and D is still caught. But notice the
number of traces: 162 with one particle, and nearly 30,000 with two.[✓](https://github.com/timvieira/testing-randomized-algorithms/blob/master/tests/test_article.py#L175)
**Pro tip:** Make sure the total mass decreases quickly, as the loop will halt
only when the total mass is zero. This requires the support to be finite. For
problems where the support is large but finite, it may be a reasonable heuristic
to stop early if the remaining mass is small, trading exactness for a bound or
an unbiased estimate (see [Which Trace Next?](#which-trace-next)).Traces are a combinatorial space, so this strategy is only practical for smaller problems. But "smaller problems" is exactly the regime where tests should live—small enough to verify exactly, but complex enough to exercise the algorithm's logic.
## Limitations
Program tracing is not a universal tool. It asks a few things of the program under test:
-
**Discrete, finite randomness.** Every random choice must have finitely many
outcomes, and every execution must terminate. A Gaussian draw, or a loop that
can run forever (rejection sampling*with* replacement, say), has infinitely
many traces, and the tracer never finishes. Sometimes you can test a discrete
skeleton of the algorithm, or discretize the continuous part; otherwise you
are back to Monte Carlo.
-
**All randomness goes through the hook.** The program must be a
deterministic function of the values`draw` returns. A stray call to the
global RNG, iteration over an unordered set, or a race between threads breaks
that assumption, and the results are meaningless.
-
**Small instances only.** The number of traces grows combinatorially: one
particle gave 162 traces, and two gave nearly 30,000. Tests have to live on
toy instances, so bugs that only appear at scale—underflow in long
sequences, a code path that only large vocabularies reach—need other
tests.
-
**Only as good as the specification.** The tracer checks the property you
wrote down. If the target is wrong, or the property is too weak (checking
only $Z$ when $Z = 1$, for instance), it will confirm the wrong thing with
perfect confidence.
-
**Exact up to floating point.** Comparisons still need a tolerance, so a bug
confined to traces whose total probability is below that tolerance goes
unnoticed.
## Takeaways
We climbed a ladder of tests, and each rung caught a bug that the one below it missed:
| Test | What it caught | Cost |
|---|---|---|
| One run, watching the error fall | B: large bias, visible after ~50 samples | one run |
| Many runs, compared to $V/M$ | C: small bias, visible after ~400 samples | a thousand runs |
| Program tracing | D: tiny bias, visible after ~15 million samples | one run per trace (196) |
| Tracing the whole SMC pipeline | bugs in how the pieces fit, like dropping the AWRS weight | one run per trace (162) |
Testing is fundamentally about **control**. Some abstractions give us more
control over a system's behavior than others, and good tests depend on choosing
abstractions that make critical behaviors observable, reproducible, and
manipulable.
The `draw` parameter in all of the algorithms above is a good example. It's a
tiny abstraction—just a function that returns a random choice—but it gives us
enormous leverage. In production, it calls the RNG. In testing, it becomes a
tracer that enumerates every possible execution. Same algorithm, radically
different observability.
A few principles I've found useful:
-
**Theory gives you test cases.** The proper weighting definition, the
unbiasedness of importance sampling, the convergence rate of Monte
Carlo—these are all*specifications* that can be mechanically checked. If you
don't understand the theory behind your algorithm, you don't know what to test.
-
**Small exact tests beat large approximate tests.** A tracer that exhaustively
checks a 5-element distribution is worth more than a Monte Carlo test with
10,000 samples. The tracer gives you a definitive yes/no; the Monte Carlo test
gives you a squiggly line that you have to interpret.
-
**Inject errors to calibrate your tests.** A test that passes on correct code
but also passes on buggy code is useless. Deliberately break the algorithm and
verify that your test catches it.
-
**Monte Carlo is the blunt instrument of last resort.** It's indispensable
when the state space is too large to enumerate, but in software—unlike the
real world—we can almost always find ways to do better. Freezing the RNG
makes your code deterministic, but it is at best a regression test and
extremely brittle.
Of course, some behaviors are only visible from "1,000 feet"—via integration tests, stress tests, or long-horizon statistical checks. Not everything can be checked with a microscope. But the more control you have over the randomness in your program, the more you can check exactly rather than approximately.
**How this post was made.** This post was written with extensive help from
[Claude Code](https://claude.ai/claude-code). The ideas, the framing, and the
final editorial decisions are mine.
Every numerical claim in the post (which variants pass, the sample sizes, the variance ratio, the trace counts) is checked by a test in the repository, and the code shown here is the code those tests run. Mistakes are mine.
Found an error or have a suggestion? Please
[open an issue](https://github.com/timvieira/testing-randomized-algorithms/issues)
on GitHub.