codecollapse

Measure how many truly different programs your code LLM samples, not just how many pass.

The problem

Teams post-training code models report pass@k, but pass@k cannot tell "ten different correct approaches" from "the same wrong idea sampled ten times with renamed variables". Text diversity metrics (BLEU self-similarity, n-gram entropy, AST hashes) count reformatted or restructured copies as different programs, so they hide mode collapse instead of exposing it. codecollapse runs every sampled program on generated inputs, groups programs by what they actually do, and reports diversity statistics with exact finite-sample estimators. It runs offline, on the standard library alone.

How it works

Input is JSONL, one problem per line:

{"problem_id": "window_sums",
 "entry_point": "window_sums",
 "reference": "def window_sums(xs, k):\n    ...",
 "input_gen": "def generate(rng, m):\n    return [([rng.randint(-9, 9) for _ in range(5)], 2) for _ in range(m)]",
 "candidates": ["def window_sums(xs, k):\n    ...", "..."],
 "labels": ["reference", "off_by_one"]}

input_gen defines generate(rng, m), returning m tuples of positional arguments; rng is a random.Random with a fixed seed, so every program sees the same inputs. labels is optional ground truth (the synthetic generator writes it); when present, the report scores recovered classes against it with the adjusted Rand index.

1. Sandboxed execution. Each program runs in its own python -S subprocess: empty environment, a fresh temp directory as cwd, resource.setrlimit on CPU seconds and address space, socket.socket replaced by a function that raises, stdout pointed at /dev/null so print cannot corrupt the result pipe, and a SIGALRM timer per input. The child writes one line per input and flushes. The parent holds a wall-clock deadline on that stream. If the child goes silent (a C-level loop, or code that catches the timer exception), the parent kills its process group, records TIMEOUT for the input in progress, and starts a fresh worker at the next input.

2. Fingerprints. Each input produces one token: a SHA-256 prefix of the canonicalised return value together with the canonicalised arguments after the call, so in-place mutation counts as behaviour; exc:<TypeName>; TIMEOUT; CRASH if the interpreter died; or load:<TypeName> if the source did not compile or define the entry point. Canonicalisation sorts dict items and set elements, prints floats to 12 significant digits, and strips at 0x... addresses from reprs.

3. Determinism check. Each program runs twice, once with random.seed and PYTHONHASHSEED from seed 1 and once from seed 2. If the two fingerprints differ, the program is flagged nondeterministic. That catches unseeded random and code that depends on set iteration order.

4. Behavioural classes. Candidates with identical fingerprints form a class. The class whose fingerprint equals the reference's is "correct".

5. Statistics are computed per problem (n candidates, c correct, class sizes n_c) and then averaged over problems:

   metric	definition
   pass@k	1 − C(n−c, k) / C(n, k) (Chen et al. 2021), exact Fraction
   E[D_k]	expected distinct behaviours among k draws without replacement, Σ_c 1 − C(n−n_c, k) / C(n, k), exact Fraction
   Hill q0, q1, q2	richness, exp(Shannon entropy), inverse Simpson over class proportions
   Chao1	bias-corrected: S_obs + (n−1)/n · f1(f1−1) / (2(f2+1)), an estimate of total behaviours, counting ones not yet sampled
   stability curve	class count using only the first m inputs, m = 1..M, plus stable_at, the smallest m after which the count stops changing
   masking ratio	distinct AST-normalised hashes / distinct behaviours

   AST normalisation parses the source, which already discards comments and layout. It then removes docstrings and alpha-renames every identifier the program binds (arguments, assignment targets, nested functions) in order of first appearance. The entry point keeps its name. So a masking ratio of 3 means that for every real behaviour, the samples contain three structurally different programs that an AST-based diversity metric would count separately.

6. Compare mode matches problems by problem_id across two checkpoints. For each metric it computes the mean per-problem delta (after − before) and a percentile bootstrap CI that resamples problems, keeping each before/after pair together.

A synthetic corpus generator (codecollapse synth) builds candidates from four hand-written reference solutions. Planted mutations change behaviour: off-by-one, flipped comparator, dropped guard clause. Semantics-preserving rewrites change only structure or text: loop ↔ list comprehension, mirrored comparisons (x > 0 → 0 < x), += expansion, guard clause → explicit else, consistent renaming, and reformatting (indent width, blank lines, comments, docstrings).

Install and usage

Requires Python 3.9+ on Linux or macOS. There are no runtime dependencies. From a checkout of this repository:

python -m pip install -e ".[dev]"      # pytest is the only dev dependency
python -m pytest -q

Worked example: pass@1 goes up, behaviours collapse

examples/ holds two synthetic checkpoints of the same four problems, each with 12 samples per problem. before.jsonl was generated with codecollapse synth -o examples/before.jsonl --collapse 0.3 --seed 1, so most samples carry a planted bug. after.jsonl was generated with --collapse 0.9 --seed 2, so almost every sample is a rewrite of the reference.

$ codecollapse measure examples/after.jsonl
problem            n  correct  behav  AST  mask  pass@1  pass@5  pass@10  E[D_1]  E[D_5]  E[D_10]  Hill1  Hill2  Chao1  stable@  nondet
----------------  --  -------  -----  ---  ----  ------  ------  -------  ------  ------  -------  -----  -----  -----  -------  ------
positive_squares  12       11      2    7  3.50   0.917   1.000    1.000   1.000   1.417    1.833   1.33   1.18   2.00     2/50       0
window_sums       12       12      1    7  7.00   1.000   1.000    1.000   1.000   1.000    1.000   1.00   1.00   1.00     1/50       0
count_in_range    12       11      2    7  3.50   0.917   1.000    1.000   1.000   1.417    1.833   1.33   1.18   2.00     1/50       0
safe_ratio        12       12      1    4  4.00   1.000   1.000    1.000   1.000   1.000    1.000   1.00   1.00   1.00     1/50       0
MEAN               -     11.5    1.5  6.2  4.50   0.958   1.000    1.000   1.000   1.208    1.417   1.17   1.09   1.50        -       0
planted-label adjusted Rand index: 1.000

Look at window_sums: 12 samples, 7 distinct normalised ASTs, 1 behaviour. An AST-hash diversity metric would score this problem at 7. Every sample does exactly the same thing. The recovered classes also match the planted labels exactly (ARI 1.000). stable@ 2/50 means the class count stopped changing after the second input, so 50 inputs were enough.

$ codecollapse compare examples/before.jsonl examples/after.jsonl
measuring examples/before.jsonl ...
measuring examples/after.jsonl ...
metric         before  after   delta            95% CI
-------------  ------  -----  ------  ----------------  -
pass@1          0.333  0.958  +0.625  [+0.437, +0.812]  *
pass@5          0.863  1.000  +0.137  [+0.036, +0.278]  *
pass@10         0.996  1.000  +0.004  [+0.000, +0.011]
E[D_1]          1.000  1.000  +0.000  [+0.000, +0.000]
E[D_5]          2.857  1.208  -1.648  [-2.196, -1.128]  *
E[D_10]         3.409  1.417  -1.992  [-2.936, -1.083]  *
hill_q0         3.500  1.500  -2.000  [-3.000, -1.000]  *
hill_q1         3.191  1.166  -2.025  [-2.739, -1.351]  *
hill_q2         3.023  1.090  -1.932  [-2.557, -1.351]  *
chao1           3.500  1.500  -2.000  [-3.000, -1.000]  *
distinct_ast    8.750  6.250  -2.500  [-3.500, -1.500]  *
masking_ratio   2.542  4.500  +1.958  [+0.667, +3.875]  *
4 paired problems, 2000 bootstrap resamples; * marks a CI that excludes zero

pass@1 nearly triples. The expected number of distinct behaviours in five samples falls from 2.9 to 1.2. Distinct ASTs fall by only 29%, so the masking ratio almost doubles: structural variety is hiding the collapse. (Four problems make a crude bootstrap. Real comparisons should use many more.)

Commands

codecollapse measure SAMPLES.jsonl [--inputs 50] [--timeout 0.5] [--memory-mb 512]
                                   [--k 1,5,10] [--input-seed 0] [--jobs N] [--json report.json]
codecollapse compare BEFORE AFTER  [same run options] [--resamples 2000] [--confidence 0.95]
                                   [--seed 0] [--json compare.json]
codecollapse synth -o corpus.jsonl [--candidates 12] [--collapse 0.6] [--seed 0]

compare accepts either samples JSONL or a report written by measure --json, so an expensive checkpoint only has to be executed once. python -m codecollapse works the same as the codecollapse script. The JSON report holds everything in the table plus, per problem, the class membership lists, each class's outcome kinds (value, exc:ValueError, TIMEOUT, …), the stability curve, and correct_curve: how many candidates match the reference on the first m inputs. The exact pass@k and E[D_k] values appear as fraction strings next to their floats. Malformed input exits with status 2 and a message naming the file and line.

Results

This project makes no performance claim. What the test suite (83 tests, python -m pytest -q, about 15 s on an Apple-silicon laptop) establishes:

• For every n from 1 to 12 and every k ≤ n, over random class structures, pass@k and E[D_k] equal brute-force enumeration over all k-subsets, compared as Fraction values (exact equality, no float tolerance).
• Planted corpora (three seeds, 4 problems × 10 candidates each) are recovered with adjusted Rand index exactly 1.0. Every planted mutation is behaviourally distinct from the reference and from the other mutations.
• Corpora made only of semantics-preserving rewrites land in a single class on every problem, with more than one AST hash, so the masking ratio is above 1.
• while True: pass is classified TIMEOUT on every input within the time budget. So is a loop that catches the timeout exception, which the parent kills, and the next input is still evaluated.
• Exceptions are classed by type, not message. os._exit is CRASH. Unseeded random is flagged nondeterministic. Seeded random.Random(x) is not.
• A candidate that differs from the reference only on input 38 of 50 has stability curve [1]*37 + [2]*13. With M = 37 it is counted as correct.
• Chao1's absolute bias shrinks at every step from n = 15 through 40 and 120 to 400 (500 seeded multinomial trials each, 25 true classes with Zipf-like weights), and it is smaller than the bias of observed richness throughout.
• The paired bootstrap 95% CI covers the known population delta in at least 93% of 400 simulated checkpoint pairs (80 problems, pass@5).

CI runs the suite on Python 3.9 (the declared minimum) and 3.13, on Ubuntu and macOS. Lint uses ruff pinned to 0.13.0.

Design notes

Fingerprint first, compare later. Each program runs once per seed on the full set of M inputs. The stability curve then comes from grouping on fingerprint prefixes, with no re-execution. That is why the inputs must be a fixed sequence from generate(rng, m) rather than being resampled per M: the curve answers "would fewer inputs have merged these classes?" about the exact inputs used. The same stored fingerprints serve compare mode through saved JSON reports.

Two layers of timeout, and a restart instead of a verdict. A SIGALRM inside the child is cheap and handles the common infinite loop without killing anything. It fails against code that swallows BaseException and against long C calls, so the parent keeps its own per-line deadline. When the parent has to kill the child, it resumes a new worker at the next input rather than marking every remaining input TIMEOUT. The cost is interpreter startup per kill. The benefit is that a program that hangs on one edge case is not merged with one that hangs on everything. Arguments after the call are part of the fingerprint for the same reason: xs.sort() and return None both return None and must not collapse into one class.

Exact arithmetic where the claim is exact. pass@k and E[D_k] are ratios of binomial coefficients, and a float implementation passes a tolerance test while quietly losing digits for large n. Keeping them as Fractions means the brute-force tests assert equality, and the per-problem means in the summary stay exact too. Hill numbers and Chao1 are inherently real-valued, so they are floats.

Limitations

• Not a security sandbox. The subprocess, rlimits and socket patch stop accidents: runaway loops, memory blowups on Linux, stray network calls through socket. They do not stop hostile code, which can re-import _socket, write files, or spawn processes. Run untrusted model output inside a container or VM as well.
• POSIX only. It relies on resource, SIGALRM and process groups. On macOS the address-space limit is not enforced by the kernel; CPU and wall-clock limits still apply.
• Behaviour equals behaviour on the generated inputs. Two programs that differ only on inputs the generator never produces share a class. The stability curve shows whether more inputs were still splitting classes, but it cannot reveal a distinguishing input that was never generated.
• Python candidates only, called as entry_point(*args). Classes, stdin/stdout programs and multi-function APIs are not supported.
• Canonicalisation is a choice. Floats are compared to 12 significant digits. 1, 1.0 and True are different outcomes. Objects without a structural encoding are compared by type name and address-stripped repr. Exceptions compare by type name only.
• AST normalisation renames per module, not per scope. Two programs that reuse a name differently in separate functions may hash apart even though a scope-aware renaming would unify them. That can only overstate AST diversity, and so the masking ratio.
• The synthetic generator knows four reference problems and the mutation and rewrite patterns they contain. It exists to validate the measurement, not to benchmark models.
• The percentile bootstrap assumes problems are exchangeable draws from a population. With few problems the intervals are too narrow.

License

MIT. See LICENSE.