Just One Line of `random.choice`? Six Detours in Quantum-State Sampling

Once a quantum circuit has finished running and produced a final-state wavefunction $\psi$, the next step is often sampling. Under the Born rule, the probability of observing the $i$-th computational basis state is $p_i=|\psi_i|^2$. The simulator's job is to draw bitstrings from this distribution and reproduce the measurement outcomes that a real quantum computer would produce, shot after shot.

In NumPy, it looks like a one-line operation:

samples = np.random.default_rng().choice(2**n, size=shots, p=p)

Getting the same operation right in TensorCircuit-NG takes six detours: it has to run on multiple computational backends; its random inputs must work with just-in-time compilation and vectorization; it must remain accurate in single precision; it must work around JAX's global precision switch; it must preserve compatibility with existing interfaces; and the new algorithm itself must also be compilable. Together, these constraints turn a one-line operation into a systems-design problem. Let's walk through them.

More Than NumPy

The same TensorCircuit-NG code can run on NumPy, TensorFlow, JAX, and PyTorch. Each library exposes a different random-sampling interface: JAX has jax.random.choice, PyTorch has torch.multinomial, and TensorFlow's tf.random.Generator has no probability-weighted choice. Adapting each one separately is possible, but the algorithm, precision, and reproducibility would then vary by backend.

A more robust approach is to return to the definition of sampling and use inverse transform sampling. Divide the interval $[0,1)$ into $N$ segments with widths $p_0,p_1,\ldots$. Draw a uniform random number $u$ between zero and one, then find the segment it lands in. The right endpoints of the segments form the cumulative distribution:

c = cumsum(p)
i = searchsorted(c, c[-1] * (1 - u))

Cumulative sum (cumsum) and searchsorted (a binary search for a position in a sorted array) are basic operations available on every backend. This lets probability_sample be implemented once in a backend-independent abstraction layer, with all four backends sharing the same algorithm.

Where Random Numbers Come From

Random-number generators naturally carry state: each draw advances their internal state. NumPy, TensorFlow, and PyTorch generators are stateful objects; JAX takes a functional approach, with randomness determined by an explicitly passed key.

Stateful generators interact awkwardly with just-in-time compilation (JIT). JIT traces a function into a computation graph, and modifications to an external generator's state inside the function are side effects. They may be frozen as constants during tracing or leak outside the graph. Automatic vectorization (vmap) has a similar problem: each task in a batch should use its own random numbers, rather than sharing one state that gets advanced repeatedly.

TensorCircuit-NG's solution is to take randomness out of the function and pass it in as an ordinary input. probability_sample accepts a status argument with shape [shots]; each element is a uniform random number in $[0,1)$, the $u$ from above. Sampling is therefore a pure function of $(p,\text{status})$. It can be JIT-compiled directly, or vmap can be applied over status to advance thousands of trajectories at once. The result depends only on $p$ and status, not on the internal implementation of any backend's generator.

The same convention runs throughout the library. Per-qubit sampling without materializing the full state uses status with shape [batch, nqubits], and Monte Carlo trajectories for noisy circuits use externally supplied uniform random numbers to select Kraus operators.

Float32 Is Not Precise Enough

TensorCircuit-NG uses complex64 by default, which means its probability vectors use float32. This is a natural choice on a GPU: it uses half the memory and runs faster. A 24-qubit complex64 state vector takes only 128 MiB, small enough to fit on any GPU.

The problem is the resolution of float32. It has 24 significant bits, so adjacent representable numbers in $[0.5,1)$ are separated by $2^{-24}\approx 6\times10^{-8}$. A probability vector for 24 qubits has $2^{24}$ entries, making the average entry about $6\times10^{-8}$ as well—on the same scale as that spacing. Inverse transform sampling can fail here in two ways.

First, the cumsum ruler has ticks that are too coarse. Once the cumulative sum exceeds 0.5, adding a probability smaller than half the spacing may get rounded away. That segment then has zero length on the ruler and can never be sampled. A slightly larger probability may round up to a full tick, amplifying its weight. In the extreme, the cumulative sum can stall altogether. On TensorFlow, sequentially accumulating $2^{25}$ equally likely entries reaches exactly 0.5 at the $2^{24}$-th entry. Each subsequent addition of $2^{-25}$ is exactly half a spacing; under round-to-even, the sum stays put, so none of the remaining entries can be sampled.

Second, and more subtly, the dart itself is discrete. A uniform float32 random number has only about $2^{23}$ possible values, so its landing points are restricted to integer multiples of $2^{-23}$. Even if the cumulative distribution were exact, each outcome's sampling probability could only be an integer multiple of $2^{-23}$. Once there are more than $2^{23}$ outcomes, large groups of them must be impossible to sample. This is a hard limit on the available information; no amount of care in the sampling algorithm can get around it.

The combined error can be quantified with the Porter–Thomas distribution, which is characteristic of random quantum circuit outputs. The total variation (TV) distance below measures how far the actual sampling distribution is from the true one. It is calculated over all possible values of the random number and excludes finite-sample statistical fluctuations:

Qubits	TV distance
16	$1.3\times10^{-3}$
20	$2.0\times10^{-2}$
24	$0.30$

At 24 qubits, the sampled distribution has already deviated from the true distribution by 30%; considering random-number discretization alone, 57% of bitstrings can never be sampled. This happens in the most commonly used sample(allow_state=True) path, at the default precision, on a state-vector size that fits on an ordinary GPU.

The same limitation affects choice-style interfaces across libraries. Consider a float32 distribution with $2^{18}$ entries: the first entry has probability 99.5%, and the remaining 0.5% is divided evenly among the others. In 200,000 draws, we would expect about 1,000 samples from the low-probability entries. In experiments, TensorFlow's choice (also based on a float32 cumulative sum) and PyTorch's torch.multinomial failed to sample any of them. jax.random.choice uses a parallel scan to compute the cumulative sum, which avoids stalling but cannot avoid random-number discretization. NumPy internally converts everything to float64 and is unaffected.

Why Not Switch Directly to Float64?

The most direct fix is to use float64 for both the cumulative sum and the random numbers. With 53 significant bits, float64 has ample precision for any state vector that can fit in memory. If users set the global precision to complex128, this approach already works without any changes.

The obstacle is JAX in complex64 mode. Whether JAX can use 64-bit floating-point values is controlled by the global jax_enable_x64 switch. When TensorCircuit-NG is set to complex64, it turns this switch off; otherwise, JAX's default floating-point and integer types would become 64-bit, quietly eating into the memory and speed savings of single precision. With the switch off, float64 arrays do not exist in JAX: converting probabilities to float64 inside a function still produces float32, silently. JAX is also the backend used most heavily with JIT and vmap.

Since local promotion to higher precision is not an option, the solution has to work within float32.

Two-Level Sampling and Interface Compatibility

Both errors above have the same root cause: too many outcomes are being considered in a single choice. Split the choice into two stages, and each stage has a much smaller range.

The method is to divide the $N$ outcomes into roughly $2^{12}$ blocks, each containing $m=N/2^{12}$ outcomes. Once the number of outcomes exceeds $2^{24}$, switch to $\sqrt{N}$ blocks instead. The first stage performs inverse transform sampling over the blocks: compute the total probability of each block, take a cumulative sum, draw the first random number, and select a block. The second stage samples within the selected block: its cumulative sum starts at zero and is measured relative to the block's total probability, so the float32 spacing is also relative to that block. Draw a second random number to locate the specific outcome. Each stage chooses among only about $2^{12}$ objects, which is well within the range that float32 can resolve reliably.

Two-level sampling needs two random numbers, which raises an interface issue: by convention, status has always contained one number per shot. A single float32 random number carries only about 23 bits of information and cannot distinguish $2^{24}$ outcomes, no matter how it is split. The second number cannot be eliminated. The interface handles this by using different status shapes:

• Pass [shots]: behavior is exactly the same as before, so existing code continues to work.
• Pass [shots, 2]: the two numbers drive the two stages separately, yielding accurate sampling in float32.
• Omit status: the function decides automatically. In float32, when the number of outcomes exceeds $2^{12}$, it generates two random numbers internally and uses two-level sampling; otherwise it generates one and uses a single stage.

In float64 mode, one random number has 52 bits, so single-stage sampling is already sufficient and behavior stays the same. With at most $2^{12}$ outcomes—that is, up to 12 qubits—the single-stage float32 error is on the order of $10^{-4}$, so sampling also remains single-stage.

Why Hand-Code the Second-Stage Binary Search?

The first stage uses the native searchsorted: every shot searches the same sorted array of length about $2^{12}$, which is exactly its intended use. The second stage is different. Each shot selects a different block, so each must search a different row, while searchsorted can handle only one sorted array at a time. To use it, we would first have to gather each shot's row and assemble a two-dimensional array of shape [shots, m]. With 65,536 shots and 24 qubits ($m=4096$), this intermediate array would occupy 1 GiB.

A Python loop over shots, with an if deciding whether to move left or right, does not work under JIT either. The computation graph's shapes and control flow must be determined during tracing; they cannot branch on the actual values of tensors.

The solution is a fixed-round, branch-free binary search:

pos = b * m
for step in [m // 2, m // 4, ..., 1]:
    pos = where(ic[pos + step - 1] < r, pos + step, pos)

Here, b is the block selected at the first stage, and pos starts at that block's first position. ic is the flattened array of within-block cumulative sums, and r is the within-block target derived from the second random number. The search finds the first position in the block for which ic >= r. At each round it asks the same question: can we jump forward by step positions? The last element we would jump over is at pos + step - 1. Because the cumulative sum is nondecreasing, if that element is less than r, every element we would skip is also less than r, so the jump is safe. step is halved each time, from $m/2$ down to 1, which is equivalent to determining the answer's binary representation one bit at a time, from the most significant bit onward. The number of rounds, $\log_2 m$, is fixed at trace time, so JIT and vmap can be nested normally.

Results

In terms of accuracy, the TV distance for the same Porter–Thomas distribution no longer grows with system size under two-level sampling:

Qubits	Single-stage float32	Two-level float32
13	$1.6\times10^{-4}$	$8\times10^{-5}$
16	$1.3\times10^{-3}$	$8\times10^{-5}$
20	$2.0\times10^{-2}$	$8\times10^{-5}$
24	$0.30$	$1.3\times10^{-4}$

At large scales, two-level sampling is actually faster. At 24 qubits, it is 2.5 to 4 times faster on TensorFlow, 1.5 to 2 times faster on PyTorch, 10% to 15% faster on JAX, and about 2.6 times faster on NumPy with many shots. One reason is that TensorFlow and PyTorch are slow at computing cumulative sums over very long one-dimensional arrays; after blocking, the row-wise cumulative sums can be parallelized.

Behind What Looks Like One Line of Code

Looking back, none of these six detours came from quantum physics itself. Supporting multiple backends means using only operations they all share. JIT and vmap require externalized randomness and static control flow. Single precision means reexamining the resolution of every cumulative sum and every random number. JAX's global precision switch rules out the simplest fix. Existing users' code requires a backward-compatible interface. Finally, the new algorithm for addressing precision problems must itself meet JIT's requirements. These constraints pull against one another; after changing any one part, we often have to check whether the rest still holds.

Users see only that sample() returns a string of bits. The reliability of a scientific computing framework depends to a large extent on whether someone has carefully thought through everything behind those bits. That is one of the important strengths of a mature, widely used quantum software framework like TensorCircuit-NG.