The Unbiased pass@k Estimator

The problem

How do you measure how good a code model is? Run the code it writes against unit tests: a program counts as correct only if it passes. But a model can write a different program each time, so one sample says little. pass@k is the probability that at least one of $k$ generated programs passes the tests. You will implement the estimator of Chen et al. (2021).

Estimating pass@k from exactly $k$ samples is noisy. So draw $n \ge k$ samples, count the $c$ that pass, and compute

\text{pass@}k = 1 - \frac{\binom{n-c}{k}}{\binom{n}{k}}

where $\binom{n}{k}$ (" $n$ choose $k$ ") is the number of ways to pick $k$ items out of $n$ . The fraction is the probability that a random group of $k$ samples is all wrong: $\binom{n-c}{k}$ groups of $k$ among the $n - c$ wrong samples, out of $\binom{n}{k}$ groups in all. One minus it is the probability that at least one is right. This estimator has no bias.

Example. $n = 10$ samples, $c = 3$ pass.

$k = 1$ : $1 - \binom{7}{1}/\binom{10}{1} = 1 - 7/10 = 0.3$ : the share of correct samples.

$k = 5$ : $1 - \binom{7}{5}/\binom{10}{5} = 1 - 21/252 = 0.917$ .

$c = 8$ , $k = 5$ : only $n - c = 2$ wrong samples, so no group of 5 can be all wrong: pass@5 $= 1$ .

Over two problems, one with $c = 3$ and one with $c = 0$ (pass@5 $= 0$ ), the mean pass@5 is $(0.917 + 0)/2 = 0.458$ .

The stable form

The binomials overflow for $n$ in the hundreds ( $\binom{200}{100}$ is about $10^{59}$ , and for larger $n$ it exceeds what a float can hold). The same number, as a product:

1 - \prod_{i = n-c+1}^{n} \left(1 - \frac{k}{i}\right)

Every factor is between 0 and 1, so nothing overflows. (If $n - c < k$ the answer is 1, as in the example; the product form would hit a negative factor.)

The biased shortcut

It is tempting to estimate $\hat p = c/n$ and compute $1 - (1 - \hat p)^k$ . It is biased: it treats the $k$ samples as independent draws with replacement.

Example. $n = 2$ , $c = 1$ , $k = 2$ : drawing both samples certainly includes the correct one, so pass@2 $= 1$ . The shortcut gives $1 - 0.5^2 = 0.75$ .

What to implement

pass_at_k(n, c, k): the estimator for one problem, as a float. If $n - c < k$ return 1.0; otherwise the product form.
mean_pass_at_k(ns, cs, k): the mean over problems (arrays of $n$ and $c$ , one entry per problem).
naive_pass_at_k(n, c, k): the biased shortcut, so that the tests can compare.

Requirements

Raise ValueError if $k > n$ .