counting: vectorize dynamic_window="prob", bit-identically

2.6x faster (0.910s -> 0.355s on 200k tokens at window=5), and -- the
part that matters more -- bit-identical to the frozen reference, so every
number ever recorded with this configuration still reproduces.

The equivalence gate asserted the opposite. Its docstring said
bit-identity was "not achievable" for randomized configurations because
"any sensible vectorization draws a whole array up front from a numpy
generator". The premise was the mistake, not the conclusion drawn from
it: iterate_tokens draws exactly one randint(1, window) per token in
token order, and a comprehension over the flattened chunk draws the same
numbers in the same order. Measured, the draws are not the expensive part
anyway -- 0.08s per 200k tokens against the ~0.9s of emission they gate.

So the gate now holds this configuration to assert_array_equal across all
four windows and both seeds, and its docstring records why the old claim
was wrong rather than quietly dropping it.

subsample="prob"/"dirty" stay on the loop, and the measurement is the
reason: they discard so many tokens that they are already the FASTEST
configurations there are -- 0.12-0.13s against 0.31s for the vectorized
deterministic default, because so few pairs survive. Vectorizing them
would optimize the cheap case and buy a second implementation to keep in
sync. Their draw streams also interleave in a way the prob window does
not (subsampling draws precede window draws, and under the clean variant
a dropped token consumes one but not the other), so it would be the
delicate work as well as the pointless kind.

Author: jfilter

CHANGELOG.md

@@ -4,6 +4,24 @@
 
 ### Added
 
+-   **`dynamic_window="prob"` is vectorized too, and is now bit-identical.**
+    2.6x faster (0.910s to 0.355s on 200k tokens at `window=5`) — and, unlike a
+    numpy-RNG rewrite would have been, it reproduces the *exact* matrices this
+    configuration produced before, so results recorded with it still hold.
+
+    The equivalence gate previously stated that bit-identity was "not achievable"
+    for any randomized configuration. That assumed the vectorization would draw
+    from a numpy generator up front; it does not have to. `iterate_tokens` draws
+    one `randint(1, window)` per token in token order, and a comprehension over
+    the flattened chunk draws the same numbers in the same order. The gate now
+    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.
+
 -   **`bench/bench_svd.py` — which SVD backend to use, measured.** The package
     offered three (`scipy` exact, `gensim` and `scikit` randomized) and never said
     which to pick. It reports speed *and* fidelity in the terms the package

hyperhyper/pair_counts.py

@@ -345,6 +345,31 @@ def uses_rng(self):
         """
         return self.subsampler_prob is not None or self.dynamic_window_prob
 
+    def is_vectorizable(self):
+        """
+        Whether `count_texts_vectorized` can handle this configuration.
+
+        Two families qualify, and what they have in common is that the *draw
+        stream* can be reproduced exactly -- either because there is none, or
+        because it is one draw per token in token order, which a list
+        comprehension over the flattened chunk gives verbatim:
+
+        * everything deterministic (no draws at all);
+        * `dynamic_window="prob"` **without** subsampling: exactly one
+          `randint(1, window)` per token, in order.
+
+        `subsample="prob"`/`"dirty"` stay on the loop. Not because their stream
+        could not be reproduced -- it could, with more care, since the
+        subsampling draws for a sentence all precede its radius draws -- but
+        because they are already the *fastest* configurations by a wide margin:
+        they discard so many tokens that few pairs survive, so vectorizing them
+        would optimize the cheap case and buy a second code path to keep in
+        sync.
+        """
+        return not self.uses_rng() or (
+            self.dynamic_window_prob and self.subsampler_prob is None
+        )
+
     def count_texts(self, texts, rng):
         """
         The counting itself, split out from `__call__` so that a caller can time
@@ -354,8 +379,8 @@ def count_texts(self, texts, rng):
         bit-identical matrix several times faster; everything else keeps the
         Python loop. See `count_texts_vectorized`.
         """
-        if not self.uses_rng():
-            vectorized = self.count_texts_vectorized(texts)
+        if self.is_vectorizable():
+            vectorized = self.count_texts_vectorized(texts, rng)
             if vectorized is not None:
                 return vectorized
 
@@ -376,19 +401,24 @@ def count_texts(self, texts, rng):
                 counter[pair[0], pair[1]] += pair[2]
         return to_count_matrix(counter, self.vocab_size)
 
-    def count_texts_vectorized(self, texts):
+    def count_texts_vectorized(self, texts, rng=None):
         """
         Count a whole chunk with numpy instead of the per-token Python loop.
 
         Returns the same CSR matrix as the loop, **bit for bit**, or `None` when
         the chunk is too large to hold its event arrays (the caller then falls
         back to the loop, which is slower but streams).
 
-        Only for configurations that draw no random numbers. The randomized
-        modes are left on the loop deliberately: their draw order per token is a
-        contract (see `iterate_tokens`), they emit far fewer pairs so they are
-        already the cheap case, and duplicating that logic would be the kind of
-        second implementation that drifts.
+        Handles the configurations `is_vectorizable` accepts: everything
+        deterministic, plus `dynamic_window="prob"` without subsampling.
+
+        The `"prob"` window keeps its exact draw stream rather than switching to
+        a numpy generator. `iterate_tokens` draws one `randint(1, window)` per
+        token in token order, and a comprehension over the flattened chunk draws
+        the same numbers in the same order -- so this stays bit-identical to the
+        loop *and* to every result recorded before it existed, which a numpy RNG
+        could never be. The draws themselves are not the expensive part
+        (~0.08s per 200k tokens against ~0.9s for the emission they gate).
 
         How bit-identity is kept
         ------------------------
@@ -428,8 +458,18 @@ def count_texts_vectorized(self, texts):
         sentence_end = sentence_start + np.repeat(lengths, lengths)
         centre = np.arange(len(flat), dtype=np.int64)
 
-        lo = np.maximum(sentence_start, centre - self.window)
-        hi = np.minimum(sentence_end, centre + self.window + 1)
+        if self.dynamic_window_prob:
+            # one draw per token, in token order -- exactly what the loop does
+            radius = np.fromiter(
+                (rng.randint(1, self.window) for _ in range(len(flat))),
+                dtype=np.int64,
+                count=len(flat),
+            )
+        else:
+            radius = self.window
+
+        lo = np.maximum(sentence_start, centre - radius)
+        hi = np.minimum(sentence_end, centre + radius + 1)
         span = hi - lo  # still includes the centre itself
         n_events = int(span.sum())
         if n_events > MAX_VECTORIZED_EVENTS:

tests/test_pair_counts_equivalence.py

@@ -32,16 +32,31 @@
 If a vectorized rewrite cannot hit bit-identity here, that is a finding to
 discuss, not a tolerance to widen.
 
-Randomized configurations -- STATISTICAL EQUIVALENCE ONLY
-----------------------------------------------------------
-    dynamic_window == "prob"   or   subsample == "prob"
-
-Bit-identity is *not achievable* here and demanding it would be a bug in the
-test, not in the code. The current implementation draws one number per token
-from `random.Random`, interleaved with the counting loop; any sensible
-vectorization draws a whole array up front from a `numpy` generator. Different
-RNG, different draw order, different numbers -- at an identical seed. The
-matrices will differ per seed no matter how correct the rewrite is.
+`dynamic_window="prob"` without subsampling -- ALSO BIT-IDENTICAL
+------------------------------------------------------------------
+This section used to say bit-identity was "not achievable" for every randomized
+configuration. That was wrong, and the reasoning behind it is worth keeping
+because it is a plausible trap: it assumed a vectorization would draw its
+numbers from a *numpy* generator up front, which would indeed give different
+numbers at the same seed no matter how correct the rewrite was.
+
+It does not have to. `iterate_tokens` draws exactly one `randint(1, window)` per
+token in token order, and a comprehension over the flattened chunk draws the
+same numbers in the same order -- so the vectorized counter reproduces the
+stream verbatim and stays bit-identical to the frozen reference, and therefore
+to every number recorded before it existed. Held to `assert_array_equal` below,
+across all four windows and both seeds.
+
+Randomized subsampling -- STATISTICAL EQUIVALENCE ONLY
+-------------------------------------------------------
+    subsample in ("prob", "dirty")
+
+These stay on the Python loop, so nothing about their output has changed. The
+statistical machinery below remains the standard they are held to, because a
+future vectorization of *them* would face the interleaving the paragraph above
+sidesteps: the subsampling draws for a sentence precede its window draws, and
+under the clean variant a dropped token consumes a subsample draw but no window
+draw.
 
 What must survive is the *distribution*. So these are checked by:
 
@@ -892,3 +907,53 @@ def test_vectorized_counter_handles_a_single_token_sentence(grid_corpus):
         counter[pair[0], pair[1]] = counter.get((pair[0], pair[1]), 0) + pair[2]
     expected = pair_counts.to_count_matrix(counter, grid_corpus.vocab.size)
     np.testing.assert_array_equal(fast.toarray(), expected.toarray())
+
+
+@pytest.mark.parametrize("seed", SEEDS)
+@pytest.mark.parametrize("window", WINDOWS)
+def test_prob_window_without_subsampling_is_bit_identical(grid_corpus, window, seed):
+    """
+    `dynamic_window="prob"` is randomized and *still* bit-identical.
+
+    The vectorized counter draws its radii from the same `random.Random`, one
+    per token in token order, rather than from a numpy generator -- so the
+    stream is the one the Python loop produced, and every number recorded before
+    the rewrite still reproduces. See the module docstring for why this was
+    previously believed impossible.
+    """
+    live, ref = _both(
+        grid_corpus,
+        window=window,
+        dynamic_window="prob",
+        subsample=None,
+        subsample_factor=SUBSAMPLE_FACTOR,
+        seed=seed,
+    )
+    np.testing.assert_array_equal(
+        live,
+        ref,
+        err_msg=(
+            f"live count_pairs diverged from the frozen reference for "
+            f"window={window} dynamic_window='prob' seed={seed}. The radius "
+            f"draws must come from random.Random in token order, so this is "
+            f"reproducible despite being randomized."
+        ),
+    )
+    assert ref.sum() > 0
+
+
+def test_prob_window_is_vectorized_but_subsampling_is_not(grid_corpus):
+    keep = pair_counts.subsample_keep_probabilities(grid_corpus.counts, 1.0)
+    assert _closure(grid_corpus, dynamic_window="prob").is_vectorizable() is True
+    # ... and adding subsampling takes it back off the fast path, because the
+    # two draw streams interleave
+    assert (
+        _closure(
+            grid_corpus, dynamic_window="prob", subsample="prob", subsampler_prob=keep
+        ).is_vectorizable()
+        is False
+    )
+    assert (
+        _closure(grid_corpus, subsample="dirty", subsampler_prob=keep).is_vectorizable()
+        is False
+    )