Ideas for Testing Randomized Algorithms

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

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 #

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 #

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 #

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, 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, 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). 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.

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, in this post's repository. It includes a traces(program) helper that runs this loop for you.

Tracer implementation #

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:

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✓—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:

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.✓ 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.✓) 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). 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). 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, 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 = '<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).

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.

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✓, 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.✓

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?).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 samplingwith 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 valuesdraw 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 allspecifications 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. 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 on GitHub.

source & further reading

timvieira.github.io — original article