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 key accounting fact is that 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 the parts 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). An exact positive integer (or None) – a fractional value raises rather than silently truncating.

• 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. An exact positive integer – a fractional or non-positive value raises.

• bin_width_bits (float) – quantization resolution of the count index (finer = exacter ordering, more memory). The defaults match the distribution count-DP. Ordering from unrank/slice (via seek_index()) is exact between fine buckets (width bin_width_bits / oversample bits) but NOT guaranteed within one – see slice(). Sequences whose log_density differs by less than one bucket’s width can surface in either relative order; this is most visible when several near-tied candidates cluster at the head of a small/short-sequence model. Raise oversample or lower bin_width_bits to shrink it. oversample is an exact positive integer – a fractional or non-positive value raises.

• oversample (int) – quantization resolution of the count index (finer = exacter ordering, more memory). The defaults match the distribution count-DP. Ordering from unrank/slice (via seek_index()) is exact between fine buckets (width bin_width_bits / oversample bits) but NOT guaranteed within one – see slice(). Sequences whose log_density differs by less than one bucket’s width can surface in either relative order; this is most visible when several near-tied candidates cluster at the head of a small/short-sequence model. Raise oversample or lower bin_width_bits to shrink it. oversample is an exact positive integer – a fractional or non-positive value raises.

• 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). Each chunk’s results must exactly match the chunk’s own size – a short or long result raises rather than silently dropping/ignoring the mismatch (MXR-080-0223).

• batch_size (int) – prefixes per batched forward. An exact positive integer – a fractional or non-positive value raises.

• 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. An exact positive integer (or None) – a fractional value raises rather than silently truncating.

• 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. Must return exactly one score per (in-support) sequence requested – a short or long result raises (MXR-080-0223).

• 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. Must return exactly one result per position – a short result raises rather than silently leaving some positions uncached (MXR-080-0223).

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 it is off the model’s DECLARED support (wrong length, missing/misplaced eos, or past the depth cap; see _in_declared_support()) or if any token along the way is locally unavailable 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. The scorer’s returned cardinality is validated (exactly one score for the one sequence requested) before it is trusted – see score_sequences()’ own note (MXR-080-0223).

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 scoring every structurally in-support sequence (one forward per sequence, all positions in parallel); a sequence off the declared support (see _in_declared_support()) scores -inf directly, without ever reaching the external scorer. The scorer’s returned cardinality is then validated – exactly one score per in-support sequence sent to it – before it is trusted (MXR-080-0223); a short or long result raises rather than silently misaligning or dropping scores. Otherwise falls back to per-sequence log_density() over the cached walk, which applies the identical structural check. 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.

The result must carry exactly one entry per position (MXR-080-0223): a short result used to be silently accepted by slicing to results[:len(seq)], leaving the un-covered tail of prefixes uncached with no signal that coverage was incomplete. That is validated – and any cache mutation withheld – before anything is written.

Parameters:

sequence (Iterable[Any])

Return type:

None

structural_fine_bucket(sequence, quantizer)

Return the quantized fine bucket for a sequence log density.

Parameters:

• sequence (Iterable[Any])

• quantizer (Quantizer)

Return type:

int

sampler(seed=None)

Return an autoregressive sampler.

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.

max_depth_bits is reconciled on every call, not just honored on the first (MXR-080-0225): a later call requesting a GENUINELY DEEPER cap raises the cached index’s own ceiling in place – SeekIndex already deepens its build lazily on demand, so widening the ceiling is enough; nothing already built is discarded or rebuilt. A later call requesting a smaller-or-equal cap is a no-op: the cached index already satisfies it, and lowering an already-built ceiling would only forbid further deepening the index could otherwise still do, for no benefit. Previously the FIRST call’s cap was silently permanent – a later, larger cap returned the same object with the old, too-small ceiling still in effect, so query results silently depended on unrelated call order.

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).

k <= 0 returns [] without starting the best-first listing at all (MXR-080-0224 – checked before any iteration, not after the loop body has already appended one item and only THEN noticed the bound was already satisfied, which used to make top_k(0) return one item instead of none).

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, QUANTIZED order) and its exact log-probability.

Random access through the count index, so the ordering carries the index’s quantization (unrank()): exact between fine buckets (width bin_width_bits / oversample bits), unspecified within one – two sequences whose log-probabilities differ by less than a bucket’s width can come back in either relative order, so rank(unrank(i)) can differ from i for near-tied neighbours. The returned log-probability is always exact regardless. Shrink the ambiguity window by raising oversample or lowering bin_width_bits; for a strictly-sorted head use top_k() (exact best-first), and for a sequence’s exact rank use rank().

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.

Read off unrank(), so it inherits the same fine-bucket quantization of the ordering (exact between buckets, unspecified within one); the returned log-probability itself is exact.

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.

Deliberately does NOT delegate to the generic count_dp_top_p(): that function assumes its exact per-item mass histogram (bucketed by the floor of the EXACT total log-density) and its count index share one bucket numbering, which holds for Composite/Record/Sequence but not here – quantized_count_index() buckets a sequence by the SUM of its per-step floor-quantized buckets, and floor(a) + floor(b) <= floor(a + b), so the structural bucket is systematically <= the exact one (same discrepancy mass_above() documents). Looking counts up at the exact-mass bucket key lands on the wrong (usually empty) slot in the count histogram – confirmed: it returns a nucleus size of 0 against a brute-force ground truth of 5. This mirrors mass_above()’s fix instead: every quantity is derived from the SAME structural count histogram, with each bucket’s per-item probability bounded by [2**(-(b + steps_bound) * bpb), 2**(-b * bpb)] (identical bound, same derivation), so the result is a provable bracket rather than an exact mass – the same contract shape count_dp_top_p documents, just derived from bounds instead of a wrongly-keyed exact histogram.

Parameters:

p (float)