mixle.enumeration.autoregressive module

Make any autoregressive model count-/threshold-/unrank-able by mixle’s enumeration machinery.

model_enumeration already lists an autoregressive model’s sequences in descending probability (best_first_decode). That is the right tool for the top handful, but it does not scale: to reach the k-th most probable sequence it must expand ~k prefixes, so a rank like 1e8 is hopeless.

This module adds the count / threshold / unrank surface for the same next_logprobs(prefix) callback, so you can answer the questions that do not require listing:

• count(min_log_prob) – how many sequences are at least this probable (without listing them),

• threshold(rank) – the log-probability of the k-th most probable sequence (the top-k boundary),

• unrank(i) – the i-th most probable sequence, by random access (one model query per step), and

• mass_above(min_log_prob) – a bracket on the cumulative probability of that head.

The trick (see notes/enumerating-a-language-model.md): the number of model forward passes is bounded by the number of distinct prefixes (<= V^(L-1)), not by the rank k. We build a count histogram per prefix and compose them up the prefix tree – but because each step p(x_t | prefix) is a distinct function of the prefix, the children are not independent, so this is a tree recursion (sum of per-token shifted child histograms), not the independent-factor convolution that convolve_indices() does for Composite.

The bridge is a thin adapter, AutoregressiveEnumerable, that implements just enough of the distribution count-index contract (quantized_count_index(), log_density(), structural_fine_bucket()) that the existing drivers – count_budget_index() and the density_rank seek/rank/cumulative/nucleus functions – work on it unchanged.

Example (transformer-style next-token decoding):

import numpy as np
def next_logprobs(prefix):
    logits = my_transformer(prefix)                 # (vocab,) -> numpy
    lp = logits - logsumexp(logits)                 # log_softmax (<= 0)
    return list(enumerate(lp))                       # [(token_id, log_prob), ...]

ar = AutoregressiveEnumerable(next_logprobs, max_len=2)   # fixed-length: support = all length-2 sequences
ar.threshold(10**8)        # log-prob of the 100,000,000-th most probable length-2 sequence
ar.count(min_log_prob)     # how many length-2 sequences are at least that probable
ar.unrank(10**6)           # the millionth most probable sequence, without listing the first 1e6
ar.top_k(5)                # the 5 most probable (exact best-first; for small k)

ar = AutoregressiveEnumerable(next_logprobs, eos=EOS)    # terminating: support = ONLY eos-terminated
ar.unrank(10**6)           # the millionth most probable COMPLETE sequence (ends in eos), of any length

Support: a fixed-length model (max_len) has support on every length-max_len sequence; a terminating model (eos) has support ONLY on eos-terminated sequences, of any length, bounded by the probability budget rather than a length cap. An un-terminated truncation has zero mass as an output and is never counted.

autoregressive_count_index(steps, prefix, depth, quantizer, max_fine_bucket, eos=None, branch_cap=None)

Tree-recursive count index over completions of prefix up to depth more tokens.

Returns (CountIndex, truncated). The histogram counts completions by fine bucket of total bits; CountIndex.get_in_bucket(fb, offset) unranks the structural (token, ...) sequence and its exact log-probability. truncated is True if any completion was dropped at the max_fine_bucket depth bound (so a caller can deepen).

Each step’s bits -log2 p(x_t | prefix) are added to every completion bucket via CountHistogram.shift(); the per-token children are pooled with CountHistogram.add(). Because steps are taken in descending probability, once a token’s own bits exceed the remaining budget every later token does too, so the loop can stop – this is what bounds the work to the live prefixes.

branch_cap recurses into only the top-cap in-budget tokens per node – the certified approximation for wide vocabularies. Each skipped token’s subtree is bounded soundly (completions within r remaining bits number at most 2**r, since their conditional probabilities sum to at most 1) and the total accumulates in CountIndex.dropped_upper: the true in-budget count lies in [total(), total() + dropped_upper]. Dropped tokens do NOT set truncated (deepening cannot recover them; raising branch_cap can).

Parameters:

• steps (Callable[[tuple], list[tuple[Any, float]]])

• prefix (tuple)

• depth (int)

• quantizer (Quantizer)

• max_fine_bucket (int)

• eos (Any)

• branch_cap (int | None)

Return type:

tuple[CountIndex, bool]

class AutoregressiveEnumerable(next_logprobs, max_len=None, eos=None, max_depth=1024, bin_width_bits=1.0, oversample=8, batch_next_logprobs=None, batch_size=256, count_mode='auto', branch_cap=None, batch_score_sequences=None, all_position_logprobs=None)

Bases: object

Adapter: an autoregressive next_logprobs(prefix) model as a count-/rank-/unrank-able object.

The support depends on the model. A fixed-length model (max_len set, no eos) has support on every length-max_len sequence. A terminating model (eos set) has support only on sequences that end in eos – so the enumeration counts/ranks/unranks exactly those, of any length, bounded by the probability budget (a tight terminating model has finitely many sequences above any threshold). A length bound is NOT the support boundary there: an un-terminated truncation has zero mass as an output and is never counted. max_depth is only a safety cap on recursion for non-tight models.

Parameters:

• next_logprobs (Callable[[tuple], Iterable[tuple[Any, float]]]) – next_logprobs(prefix) -> [(token, log_prob), ...] – the next-token log-probabilities (<= 0) given a prefix tuple, e.g. the log_softmax of a transformer’s next-token logits. For speed it may instead return the (tokens, log_probs) numpy-array pair (skips per-token boxing). For a terminating model eos must be one of the tokens it can return.

• max_len (int | None) – for a fixed-length model, the sequence length (the support is all length-max_len sequences). Omit for a terminating model (or pass it as a hard length cap, but un-terminated truncations are still dropped).

• eos (Any) – end-of-sequence token. When given, the model is terminating: only sequences ending in eos are in the support.

• max_depth (int) – safety bound on recursion depth for a terminating model (the probability budget is the real bound). Raise it if a tight model legitimately produces very long sequences.

• bin_width_bits (float) – quantization resolution of the count index (finer = exacter ordering, more memory). The defaults match the distribution count-DP.

• oversample (int) – quantization resolution of the count index (finer = exacter ordering, more memory). The defaults match the distribution count-DP.

• batch_next_logprobs (Callable[[list[tuple]], list[Any]] | None) – optional batch_next_logprobs([prefix, ...]) -> [result, ...] scoring many prefixes in one (padded) forward. When given, the count index warms its forward cache breadth-first in batch_size chunks – the large speed-up for transformers, where one-at-a-time forwards dominate (e.g. distilGPT-2 length-2 to rank 1e5: ~25 s one-at-a-time -> ~1 s batched).

• batch_size (int) – prefixes per batched forward.

• count_mode (str) – how counts are carried past the int64-exact regime (budgets over ~2**60 sequences). 'auto' (default) switches the numpy fast path to float64 there – approximate counts (exact below 2**53, ~1e-16 relative error per pooling beyond) at full numpy speed. 'exact' preserves arbitrary-precision counts by falling back to the slow Python recursion. 'float' forces float64 everywhere.

• branch_cap (int | None) – recurse into only the top-branch_cap in-budget tokens per prefix – the certified approximation for wide (LLM-sized) vocabularies, shrinking the tree by ~V/cap per level. The skipped remainder is soundly bounded (count_bracket/dropped_upper: a skipped subtree with r remaining budget bits holds at most 2**r completions); enumeration covers the sub-support of sequences whose every token is among its context’s top-branch_cap.

• batch_score_sequences (Callable[[list[tuple]], Any] | None) – optional teacher-forcing scorer [sequence, ...] -> array of total log-probs – ONE forward per sequence (all positions score in parallel) instead of one forward per token. Used by score_sequences() and, when a sequence’s prefixes are not already cached, by log_density(); the substrate for draft-rescored (speculative) enumeration.

• all_position_logprobs (Callable[[tuple], list[Any]] | None) – optional sequence -> [next_logprobs result for seq[:d], d in 0..len-1] – one forward yields the full next-token distribution at EVERY position; harvested into the forward cache by harvest(). Makes corpus-calibrated envelopes ~L-times cheaper.

The model is queried lazily and memoized by prefix, so deepening the index (or recomputing a log-density) never re-runs a forward pass it has already seen. With integer tokens the histogram build is the numpy fast path (int64 counts below ~``2**60`` budgets; float64 beyond, per count_mode).

quantized_count_index(quantizer, max_fine_bucket)

Count index over the model’s support (length-max_len sequences, or all eos-terminated sequences), bounded by the bit budget max_fine_bucket.

Parameters:

• quantizer (Quantizer)

• max_fine_bucket (int)

Return type:

tuple[CountIndex, bool]

log_density(sequence)

Exact total log-probability of a sequence (-inf if any token is off-support given its prefix).

When a batch_score_sequences scorer is configured and any of the sequence’s prefixes is not already cached, the score comes from ONE teacher-forcing forward instead of one forward per token.

Parameters:

sequence (Iterable[Any])

Return type:

float

score_sequences(sequences)

Exact total log-probabilities of many sequences – batched teacher forcing when available.

With batch_score_sequences this is one call (one forward per sequence, all positions in parallel); otherwise it falls back to per-sequence log_density() over the cached walk. The rescoring primitive for draft-based (speculative) enumeration.

Parameters:

sequences (list[Any])

Return type:

ndarray

harvest(sequence)

Cache the next-token distribution at every prefix of sequence from one forward.

Requires all_position_logprobs; a no-op without it. Feeding typical sequences (a corpus, a provider’s fast generations) through this warms the same memo cache the count index and the envelope read – L cache entries per model call.

Parameters:

sequence (Iterable[Any])

Return type:

None

structural_fine_bucket(sequence, quantizer)

Parameters:

• sequence (Iterable[Any])

• quantizer (Quantizer)

Return type:

int

sampler(seed=None)

Parameters:

seed (int | None)

Return type:

_ARSampler

enumerator()

A descending-probability iterator over the support (lazy best-first); use top_k for the head.

seek_index(*, max_depth_bits=4096.0)

The cached persistent SeekIndex over this model.

Built lazily on first use and reused by every convenience query (unrank / count / threshold / mass_above), deepening in place when a query needs more depth – so a sweep of a thousand unranks pays for one tree build, not a thousand. The forward cache is shared with it, so deepening only runs new forwards for newly-live prefixes.

Parameters:

max_depth_bits (float)

budget_index(budget_bits, max_depth_bits=4096.0)

The count-budget seek index covering at least 2**budget_bits sequences (for unrank/iterate).

Parameters:

• budget_bits (float)

• max_depth_bits (float)

envelope_index(*, n_paths=64, seed=0, budget_bits=64.0)

An AREnvelopeIndex over this model – approximate enumeration at depths the exact tree index cannot reach (O(L) forwards per unrank instead of Theta(count) tree expansion; exact for iid-step models, mean-field estimate otherwise).

Parameters:

• n_paths (int)

• seed (int)

• budget_bits (float)

top_k(k)

The k most probable sequences, exact, by best-first listing (use for small k).

Parameters:

k (int)

Return type:

list[tuple[tuple, float]]

count(min_log_prob)

How many sequences have log_density >= min_log_prob – computed from counts, not listed.

With branch_cap set this is the count over the capped sub-support (a sound lower bound); count_bracket() adds the certified upper bound including the skipped remainder.

Parameters:

min_log_prob (float)

Return type:

int | float

count_bracket(min_log_prob)

A sound [lo, hi] bracket on the number of sequences with log_density >= min_log_prob.

lo counts the (exactly enumerated) kept sub-support; hi adds dropped_upper – the certified bound on completions excluded by branch_cap (identical to lo when no cap is set).

Parameters:

min_log_prob (float)

Return type:

tuple[float, float]

unrank(i)

The i-th most probable sequence (0-based) and its exact log-probability, by random access.

Parameters:

i (int)

Return type:

tuple[tuple, float]

threshold(rank)

Log-probability of the rank-th most probable sequence – the boundary of the top-rank set.

Parameters:

rank (int)

Return type:

float

mass_above(min_log_prob)

A (lower, upper) bracket on the total probability of sequences with log_density >= min_log_prob.

Computed from the count histogram alone (no enumeration): each fine bucket of c sequences contributes between c * 2**(-hi_bits) and c * 2**(-lo_bits), where the bucket spans [lo_bits, hi_bits) of information. Tighten by raising oversample.

Parameters:

min_log_prob (float)

Return type:

tuple[float, float]

seek(index)

CountDPSeekResult at descending index (with a bracket).

Parameters:

index (int)

rank(sequence)

DensityRankResult – rank + cumulative mass of a sequence.

Parameters:

sequence (Iterable[Any])

cumulative(sequence)

G(seq) = P(p(Y) >= p(seq)) – total mass of sequences at least as probable as seq.

Parameters:

sequence (Iterable[Any])

nucleus_size(p)

Size of the minimal >= p-mass set (CountDPTopPResult), without materializing it.

Parameters:

p (float)