counting: vectorize the subsampling paths, bit-identically

`subsample="prob"`/`"dirty"` were the last two configurations still on the
Python loop, on the stated grounds that they were already the fastest ones and
vectorizing them "would optimize the cheap case". That was measured and it was
wrong in the usual way: they were fastest per surviving pair and the *slowest
per call*, which is the number a user waits on. At window=5 on 720k tokens they
were the two slowest rows in the benchmark.

The real obstacle was the one the equivalence gate named: the two draw streams
interleave. Within a sentence `iterate_tokens` draws one random() per
subsample-eligible token, and only then, knowing which survived, one
randint(1, window) per survivor -- so neither stream can be drawn in a single
pass over the chunk. `_subsample_draws` walks that interleaving sentence by
sentence in Python and leaves only the emission to numpy, which is the right
split: the draws were the cheap half all along.

    subsample="prob"     1.19s -> 0.39s   (3.0x)
    subsample="dirty"    1.39s -> 0.42s   (3.3x)
    both knobs random    1.04s -> 0.50s   (2.1x)

Because the streams are reproduced rather than replaced, `subsample="prob"` is
now bit-identical to the frozen f68cc74 reference. The equivalence grid grows
from 48 exact cells to 96 -- all four windows x all four dynamic_window modes x
three subsample modes x both seeds -- and the statistical comparison for "prob"
is removed: bit-identity implies equality of every moment, so keeping it would
cost ten runs per cell to prove less. "dirty" keeps its statistical check (the
frozen reference predates the mode) and gains an exact one against our own loop.

Fixes a latent bug found by extending that fallback test to the randomized
modes: a chunk above MAX_VECTORIZED_EVENTS had already drawn from the RNG before
discovering it was too large, so the loop resumed from an advanced generator and
a corpus counted differently depending only on how it was chunked. count_texts
now snapshots and restores the RNG state around the attempt.

Also drops `is_vectorizable`, which now has no false case, and corrects the
benchmark's note that the "prob" rows' checksums are expected to move -- nothing
vectorized the RNG in the end, so every row is bit-reproducible.

Author: jfilter

CHANGELOG.md

@@ -44,10 +44,47 @@
     holds this configuration to `assert_array_equal` across every window and
     seed, and the docstring records why the old claim was wrong.
 
-    `subsample="prob"`/`"dirty"` deliberately stay on the Python loop: they
-    discard so many tokens that they are already the *fastest* configurations
-    (0.12-0.13s against 0.31s for the vectorized default), so vectorizing them
-    would optimize the cheap case and buy a second code path to keep in sync.
+-   **`subsample="prob"`/`"dirty"` are vectorized too, and are also
+    bit-identical.** ~3x faster (1.19s to 0.39s for clean, 1.39s to 0.42s for
+    dirty, on 720k tokens at `window=5`); with a randomized window on top of the
+    subsampling, 2.1x. Every configuration now takes the vectorized path.
+
+    This entry supersedes the paragraph above, which said these two "deliberately
+    stay on the Python loop" because they were already the fastest configurations
+    and vectorizing them "would optimize the cheap case". That was measured, and
+    it was wrong in the way benchmarks usually are: they were fastest *per
+    surviving pair* and the slowest *per call*, which is the number a user waits
+    on. At `window=5` they were the two slowest rows in the table.
+
+    The real obstacle was the one the equivalence gate named: the two draw
+    streams **interleave**. Within a sentence, `iterate_tokens` draws one
+    `random()` per subsample-eligible token, and only then, knowing which tokens
+    survived, one `randint(1, window)` per survivor — so neither stream can be
+    drawn in a single pass over the chunk. `_subsample_draws` walks that
+    interleaving sentence by sentence, in Python, and leaves only the *emission*
+    to numpy. That is the right split: the draws were the cheap half all along.
+
+    Because the streams are reproduced rather than replaced, `subsample="prob"`
+    is now held to `assert_array_equal` against the frozen f68cc74 reference —
+    across all 4 windows, all 4 `dynamic_window` modes and both seeds, both knobs
+    randomized at once included. The equivalence grid grew from 48 exact cells to
+    96, and the statistical comparison for `"prob"` was **removed**: bit-identity
+    implies equality of every moment, so keeping it would have cost 10 runs per
+    cell to prove less. `"dirty"` keeps its statistical check (the frozen
+    reference predates the mode and cannot produce it) and gains an exact one
+    against this package's own loop.
+
+### Fixed
+
+-   **An oversized chunk no longer changes randomized results.** A chunk above
+    `MAX_VECTORIZED_EVENTS` falls back to the Python loop — but the vectorized
+    path had already drawn from the RNG before discovering it was too large, so
+    the loop resumed from an advanced generator. A corpus would have counted
+    differently depending only on how it happened to be chunked, which is exactly
+    what that memory cap must never do. `count_texts` now snapshots and restores
+    the RNG state around the attempt. Present since `dynamic_window="prob"` was
+    vectorized (unreleased), and found by extending the fallback test to the
+    randomized modes.
 
 -   **`bench/bench_svd.py` — which SVD backend to use, measured.** The package
     offered three (`scipy` exact, `gensim` and `scikit` randomized) and never said

README.md

@@ -185,11 +185,23 @@ tests and respect the correctness gate around the counting code.
 
 ## Future Work / TODO
 
--   Vectorize pair counting. The counting in `hyperhyper/pair_counts.py` is still
-    pure Python (`iterate_tokens`), and rewriting it in numpy is the obvious next
-    optimisation. Measured on a 1.5M-token synthetic corpus (`window=2`,
-    `dynamic_window="deter"`, 10-core M1 Pro, Python 3.12), a full
-    corpus→count→PMI→SVD→evaluate run splits as:
+-   ~~Vectorize pair counting.~~ **Done** — every configuration now goes through
+    `CountPairsClosure.count_texts_vectorized`, and the Python loop
+    (`iterate_tokens`) survives only as the streaming fallback for chunks above
+    `MAX_VECTORIZED_EVENTS` and as the readable definition of what is being
+    computed. Measured speedups at `window=5` on 720k tokens: 2.6x for
+    `dynamic_window="prob"`, ~3x for `subsample="prob"`/`"dirty"`, 2.1x with both
+    randomized at once.
+
+    Every one of those is **bit-identical** to the pre-vectorization output,
+    randomized modes included: the fast path reproduces `random.Random`'s draw
+    stream instead of replacing it with a numpy generator, so results recorded
+    before the rewrite still reproduce exactly. `bench/bench_pair_counts.py`
+    measures it, `tests/test_pair_counts_equivalence.py` enforces it.
+
+    The measurement that motivated the work, kept for scale. On a 1.5M-token
+    synthetic corpus (`window=2`, `dynamic_window="deter"`, 10-core M1 Pro,
+    Python 3.12), a full corpus→count→PMI→SVD→evaluate run split as:
 
     | stage | time | share |
     |---|---:|---:|
@@ -199,10 +211,9 @@ tests and respect the correctness gate around the counting code.
     | PMI | 0.37s | 3.0% |
     | evaluation | 0.03s | 0.2% |
 
-    Inside counting, `cProfile` puts ~55% in `iterate_tokens` and ~32% in the
-    dictionary accumulation around it — i.e. roughly 85% is the pure-Python part
-    a vectorized rewrite would replace. So the ceiling here is real: several
-    percent of end-to-end per doubling of the loop's speed.
+    Inside counting, `cProfile` put ~55% in `iterate_tokens` and ~32% in the
+    dictionary accumulation around it — i.e. roughly 85% was the pure-Python part
+    the rewrite replaced. The ceiling was real, and so was the payoff.
 
     Note that the pool usually does *not* run at this size — `count_pairs`
     estimates the serial cost and skips the pool when its ~3s spawn startup would
@@ -213,10 +224,13 @@ tests and respect the correctness gate around the counting code.
     with the loop "roughly a quarter" of that. The measurement above does not
     support it and supersedes it.)
 
-    A correctness gate is in place for whoever does it:
-    `tests/test_pair_counts_equivalence.py` compares the live counter against a
-    frozen reference (`bench/reference.py`) and requires bit-identical output on
-    the deterministic configurations. See [CONTRIBUTING.md](./CONTRIBUTING.md).
+    The correctness gate stays: `tests/test_pair_counts_equivalence.py` compares
+    the live counter against a frozen reference (`bench/reference.py`) and
+    requires bit-identical output on every configuration that reference can
+    produce. See [CONTRIBUTING.md](./CONTRIBUTING.md).
+
+-   Tokenization is now the largest single stage. It is the next thing worth
+    attacking, and unlike counting it has no correctness gate yet.
 
 ## `hyperhyper`?
 

bench/bench_pair_counts.py

@@ -57,9 +57,12 @@
      cheap smoke test; the real gate is
      `pytest tests/test_pair_counts_equivalence.py -m slow`, which compares
      whole matrices against the frozen snapshot in `bench/reference.py`.
-  2. The two `prob` rows will differ in `checksum` even when the rewrite is
-     correct -- vectorizing the RNG changes the draw order. Only their timings
-     are comparable.
+  2. So must the `prob` rows'. This note used to say the opposite -- that they
+     "will differ even when the rewrite is correct, because vectorizing the RNG
+     changes the draw order". Nothing vectorized the RNG in the end: the
+     randomized paths reproduce their draw streams from `random.Random` and only
+     vectorize the emission, so every row in this table is bit-reproducible and
+     a moved checksum is a bug on any of them.
   3. `best` (the fastest of `repeats` runs) and `tok/s` are the numbers to
      quote. `mean` is there to show how noisy the machine is; if `mean - best`
      is a large fraction of `best`, close things and run it again.
@@ -106,6 +109,11 @@
     ("w10 dyn=deter", {"window": 10, "dynamic_window": "deter", "subsample": None}),
     ("w5 dyn=prob", {"window": 5, "dynamic_window": "prob", "subsample": None}),
     ("w5 sub=prob", {"window": 5, "dynamic_window": None, "subsample": "prob"}),
+    ("w5 sub=dirty", {"window": 5, "dynamic_window": None, "subsample": "dirty"}),
+    # both knobs randomized: the case whose two draw streams interleave, and the
+    # one where the vectorized counter has the least room to win, because it has
+    # to walk that interleaving in Python
+    ("w5 dyn+sub=prob", {"window": 5, "dynamic_window": "prob", "subsample": "prob"}),
 ]
 
 # large enough that subsampling actually drops tokens instead of erasing the
@@ -261,8 +269,8 @@ def main(argv):
             f"best {min(e2e):.3f}s over {len(corpus.texts)} chunks"
         )
         print()
-        print("checksum/sum/nnz must not move on the deterministic rows; the two")
-        print("'prob' rows are expected to move once the RNG is vectorized.")
+        print("checksum/sum/nnz must not move on ANY row: the randomized modes")
+        print("reproduce their draw streams rather than replacing them.")
         print()