mixle.enumeration package

The enumeration concern — one home for “what can be enumerated, and how”.

Enumeration is not a property of distributions; it is a capability shared by distributions, relations, quantized objects, and any combinatorial model. Anything that can iterate its support in descending probability implements DistributionEnumerator (the contract) and reports the Enumerable capability; anything finite-or-structural can additionally be unranked by integer rank (RankableByIndex) through the count-budget seek index.

This module gathers that concern in one place — the contract, the capability lens, the k-best algorithms, the count-budget unranking, and the count semiring — so you can read off what enumeration is and who participates without hunting through utils/stats. The implementations live here, split by concern: generic stream primitives in streams, best-first / product search in best_first, the quantized seek/unrank index and the structural count-budget index in quantization, and the k-best combinatorial enumerators in assignment / spanning. Only the contract (DistributionEnumerator) lives in the compute layer and is re-exported here. mixle.enumeration.algorithms remains as a back-compat re-export of the stream / best-first / index names. Layout per docs/ARCHITECTURE.md.

Who plugs in: every finite/countable leaf (Categorical, Poisson, …), the combinators over enumerable children (Sequence, Composite, Mixture, …), the graph/ranking families (Markov chains, Mallows, Chow-Liu trees, spanning trees), and mixle.relations.Relation — all via the same enumerator().

class Enumerable[source]

Bases: PredicateCapability

Its support can be iterated in descending-probability order (enumerator() works).

classmethod check(obj)[source]
Parameters:

obj (Any)

Return type:

bool

class FiniteSupport[source]

Bases: PredicateCapability

Has a finite number of distinct support points (support_size() is an int).

classmethod check(obj)[source]
Parameters:

obj (Any)

Return type:

bool

class RankableByIndex[source]

Bases: PredicateCapability

Supports random access by integer rank (unranking) via the count-budget seek index.

Finite-support enumerables are rankable by index through the count-DP semiring; the implication edge finite ==> enumerable ==> rankable is exactly this composition. (Decomposable combinators over infinite-but-countable children are also rankable structurally; this conservative check covers the finite-leaf case used by the dispatch layers.)

classmethod check(obj)[source]
Parameters:

obj (Any)

Return type:

bool

supports(obj, capability)[source]

Return whether obj provides capability (a Protocol or a PredicateCapability).

Parameters:
Return type:

bool

top_k(dist, k)[source]

Return the k highest-probability (value, log_prob) pairs.

Dispatches on capability, not class: a RankableByIndex distribution could random-access by rank, but every Enumerable distribution can stream its support in descending probability — so one implementation covers all of them, including user-defined ones. Raises a clear CapabilityError for distributions that support neither.

Parameters:
Return type:

list[tuple[Any, float]]

class DistributionEnumerator(dist)[source]

Bases: ABC

Lazy iterator over the support of dist in non-increasing probability order.

Yields (value, log_prob) pairs, possibly infinitely many. Contract:

  • Each support value is yielded exactly once (deduplication is the enumerator’s responsibility).

  • log_prob equals dist.log_density(value) up to float round-off (~1e-10), and the sequence of log_probs is non-increasing up to the same tolerance.

  • Values with zero probability are skipped, never yielded.

  • Ties are broken deterministically by insertion order; no further guarantee.

Parameters:

dist (SequenceEncodableProbabilityDistribution)

top_k(k)[source]

Return the k most probable (value, log_prob) pairs (fewer if the support is smaller).

Parameters:

k (int)

Return type:

list[tuple[Any, float]]

top_p(p, max_items=None)[source]

Return the smallest descending-probability prefix whose total probability reaches p.

The nucleus / minimal high-probability set: because values are yielded in non-increasing probability order, the returned prefix is a minimum-size set of outcomes whose summed mass is >= p (e.g. p=0.95 gives a 95%-coverage support set – the discrete analogue of nucleus sampling). Accumulation stops as soon as the cumulative probability reaches p.

max_items caps how many values are pulled so an infinite or heavy-tailed support cannot run away; if the cap is hit before the threshold, the (sub-threshold) prefix gathered so far is returned. p >= 1.0 on an infinite support therefore requires max_items.

Parameters:

  • p (float) – Target cumulative probability; p <= 0 returns the empty set.

  • max_items (Optional[int]) – Hard cap on the number of values pulled.

Returns:

List of (value, log_prob) pairs in non-increasing probability order.

Return type:

list[tuple[Any, float]]

quantized_index(max_bits, bin_width_bits=1.0)[source]

Precompute a bounded bit-quantized index over this enumeration.

The index groups values by floor((-log2 p(x)) / bin_width_bits), includes only values with -log2 p(x) <= max_bits, and returns exact log probabilities for indexed values. Building the index consumes this enumerator.

Parameters:

  • max_bits (float) – Maximum information content in bits to index.

  • bin_width_bits (float) – Width of each quantized probability bin in bits.

Returns:

mixle.enumeration.algorithms.QuantizedEnumerationIndex.

rank(value)[source]

Rank and cumulative probability of value in the descending-probability order.

Returns a DensityRankResult with .rank (0-based count of strictly-more-probable outcomes; None if only the sampling estimate was used) and .cumulative_probability (G(value) = P(p(Y) >= p(value))). Exact head enumeration where the support is countable, with a Monte-Carlo fallback for the deep tail.

Parameters:

value (Any)

seek(index)[source]

The value at descending-probability index (0-based) – the inverse of rank().

Returns a CountDPSeekResult carrying the value and a provable [rank_lower, rank_upper] bracket. Uses the structural count-DP for decomposable families, so arbitrarily deep indices are reachable without enumerating the prefix.

Parameters:

index (int)

seek_certified(index)[source]

The value at descending index with a GUARANTEED bracket on its TRUE marginal rank.

Unlike seek() – whose bracket bounds only the tropical rank for a marginal family (mixture/HMM) – this widens the rank window by the family’s tropical_displacement_bits and divides out the component over-count, so the returned MarginalSeekResult [true_rank_lower, true_rank_upper] provably contains #{u : log p(u) > log p(value)}. It pins the rank exactly (.exact) for decomposable / provably-disjoint families and for shallow indices, and otherwise returns the honest provable envelope. For a decomposable family it agrees with seek().

Parameters:

index (int)

cumulative(value)[source]

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

Parameters:

value (Any)

nucleus_size(p)[source]

Size of top_p() (the minimal >= p-mass set) WITHOUT materializing it.

Returns a CountDPTopPResult with a provable size bracket, from the structural count-DP – usable when the nucleus is far too large to list.

Parameters:

p (float)

from_index(start, stop=None)[source]

Iterate (value, log_prob) in descending-probability order starting at structural start.

Yields the same stream as iterating a fresh enumerator but beginning at index start (and ending before stop if given). A fresh underlying enumeration is used, so this does not consume self. (Decomposable families admit a direct structural jump via the count-budget index; the current implementation skips the best-first prefix – the structural fast path is a WS-3 performance follow-up.)

Parameters:

  • start (int)

  • stop (int | None)

exception EnumerationError(dist, path='', reason='')[source]

Bases: NotImplementedError

Raised when a distribution (or a child of a combinator) cannot enumerate its support.

The path argument identifies the offending child within a combinator, e.g. ‘CompositeDistribution.dists[1]’.

Parameters:
Return type:

None

child_enumerator(child, path)[source]

Construct child.enumerator(), annotating EnumerationError with the child’s path.

Combinator enumerators use this so a failure deep in a nested model reports the full path to the offending leaf, e.g. ‘CompositeDistribution.dists[1] -> MixtureDistribution.components[0] -> GaussianDistribution …’.

Parameters:

  • child (ProbabilityDistribution)

  • path (str)

Return type:

DistributionEnumerator

supports_enumeration(dist)[source]

Return True if dist.enumerator() can be constructed.

Return type:

bool

class DensityRankResult(cumulative_probability, rank, exact, stderr, log_prob, method, rank_stderr=0.0)[source]

Bases: object

Outcome of a rank / cumulative-probability query.

Parameters:
cumulative_probability

G(x) = sum_{y: p(y) >= p(x)} p(y) (the descending-order CDF at x).

Type:

float

rank

number of observations strictly more probable than x (0-based position). Exact when the head enumeration resolved it; in "sampling" mode it is the rounded unbiased Monte-Carlo estimate (see rank_stderr); None only when log p(x) = -inf or an exact-analytic CDF (no count) was used.

Type:

int | None

exact

True when the head enumeration resolved the query exactly; False for a sampling estimate.

Type:

bool

stderr

standard error of cumulative_probability (0.0 when exact).

Type:

float

log_prob

log p(x).

Type:

float

method

"exact-head", "exact-exhausted", "exact-analytic", or "sampling".

Type:

str

rank_stderr

standard error of the Monte-Carlo rank estimate (0.0 when exact). Large relative to rank exactly in the deep-tail / high-entropy band where exact marginal rank is provably hard – treat the estimate as unreliable there.

Type:

float

cumulative_probability: float
rank: int | None
exact: bool
stderr: float
log_prob: float
method: str
rank_stderr: float = 0.0
density_rank(dist, value, max_exact=100_000, n_samples=20_000, seed=0, tol=1.0e-9)[source]

Rank and cumulative probability of value under dist’s descending-probability order.

Strategy:

  1. If dist supports enumeration, walk the exact descending stream, accumulating the mass of every item at least as probable as value and counting those strictly more probable. If the stream drops below p(value) (or is exhausted) within max_exact items, the rank and cumulative probability are returned EXACTLY.

  2. Otherwise (value is deeper than max_exact, or enumeration is unsupported), estimate the cumulative probability by Monte Carlo: G_hat = mean_i 1[log p(Y_i) >= log p(value)] with Y_i ~ dist. Reliable here precisely because G is large in the tail.

Parameters:

  • dist (Any) – A distribution exposing log_density and sampler; optionally enumerator.

  • value (Any) – The observation to locate.

  • max_exact (int) – Cap on items pulled from the exact enumerator before falling back to sampling.

  • n_samples (int) – Monte-Carlo sample count for the fallback.

  • seed (int) – Sampler seed for the fallback (reproducible).

  • tol (float) – Log-probability tolerance for the >= comparison (ties).

Returns:

DensityRankResult.

Return type:

DensityRankResult

count_budget_index(dist, budget_bits, bin_width_bits=1.0, oversample=8, max_depth_bits=4096.0, num_workers=None, count_mode='exact')[source]

Driver for the count-budget mode: deepen until the budget is covered, then build the index.

Calls dist.quantized_count_index(quantizer, max_fine_bucket) at geometrically increasing depths until the structural count reaches 2**budget_bits or the support is exhausted (no further truncation), then wraps the result with build_budget_index(). The exact log density of each unranked value is reported via dist.log_density.

When num_workers is greater than 1, the heavy count-histogram convolutions are computed on a process pool (parallel quantization); the parallel result is identical to the serial one. count_mode='float' carries counts as float64 (C-speed convolutions, exact below 2**53, ~1e-16 relative error per operation beyond) – the approximate-counting mode for deep budgets.

Parameters:

  • budget_bits (float)

  • bin_width_bits (float)

  • oversample (int)

  • max_depth_bits (float)

  • num_workers (int | None)

  • count_mode (str)

class SeekIndex(model, *, bin_width_bits=1.0, oversample=8, count_mode='exact', max_depth_bits=4096.0)[source]

Bases: object

Build a model’s count-budget index once and answer many enumeration queries against it.

model is anything implementing the count-index contract – every mixle distribution with a quantized_count_index (Composite/Sequence/Markov exactly; Mixture/HMM as the documented tropical bound) and AutoregressiveEnumerable (so LLMs plug in).

Queries auto-deepen: asking for a rank (or a log-prob threshold) beyond the built depth rebuilds the DP at a geometrically larger budget, reusing whatever the model itself memoizes. max_depth_bits caps the deepening; past it, queries raise IndexError for out-of-range ranks.

Parameters:

  • model (Any)

  • bin_width_bits (float)

  • oversample (int)

  • count_mode (str)

  • max_depth_bits (float)

property built_bits: float

The depth (bit budget) the index is currently built to (0 before the first build).

property truncated: bool

True when values beyond the built depth exist (deepening can still grow the index).

property builds: int

How many times the structural DP has been (re)built – the cost a persistent index amortizes.

property dropped_upper: float

Certified upper bound on in-budget values excluded by an approximation knob (e.g. branch_cap).

0.0 for exhaustive indices. The true in-budget count lies in [len(self), len(self) + dropped_upper]; deepening does not recover these (that is what distinguishes it from truncated).

ensure_bits(depth_bits)[source]

Make the index cover at least depth_bits of information (build or deepen if needed).

Parameters:

depth_bits (float)

Return type:

SeekIndex

ensure_count(n)[source]

Make the index hold at least n values (geometric deepening from the current depth).

Every rebuild at least doubles the built depth, so an ascending sequence of queries costs O(log) rebuilds total and the final (deepest) build dominates the work – the amortization a persistent index exists for.

Parameters:

n (int)

Return type:

SeekIndex

unrank(i)[source]

The i-th most probable value (0-based, quantized order) and its exact log-probability.

Parameters:

i (int)

Return type:

tuple[Any, float]

slice(start, k)[source]

Up to k values starting at rank start (one deepen at most, then table reads).

Parameters:
Return type:

list[tuple[Any, float]]

iter_from(start=0)[source]

Iterate values from rank start through the end of the built index (no auto-deepen).

Parameters:

start (int)

Return type:

Iterator[tuple[Any, float]]

count(min_log_prob)[source]

How many values have log_density >= min_log_prob – read off the fine histogram.

Exact-carrier families return an int; count_mode='float' (or a float-carrying model like the deep-budget autoregressive fast path) returns a float with the documented relative error. Mixture/HMM counts are the structural (tropical) upper bound those families document.

Parameters:

min_log_prob (float)

Return type:

int | float

threshold(rank)[source]

Log-probability of the rank-th most probable value (the top-rank boundary).

Parameters:

rank (int)

Return type:

float

rank_bracket(value)[source]

A [lo, hi] bracket on value’s quantized rank, from its structural fine bucket.

lo counts every value in strictly shallower buckets; hi adds the rest of the value’s own bucket. For Mixture/HMM the bucket is the documented tropical projection, so the bracket carries that projection’s semantics (see the family docstrings).

Parameters:

value (Any)

Return type:

tuple[int, int]

fine_histogram(min_depth_bits=None)[source]

The built fine CountIndex (ensuring at least min_depth_bits) – for mass/rank math.

Parameters:

min_depth_bits (float | None)

logit_error_bucket_slack(eps_nats, steps, quantizer)[source]

Sound bucket-shift bound when step log-probabilities carry up to eps_nats of error each.

The license to run enumeration forwards on a quantized model (int8/int4 inference, a distilled twin): enumeration never needs logits more accurate than the fine-bucket width. If every step log-probability is within eps_nats of the true model’s, a steps-long value’s accumulated fine bucket shifts by at most ceil(steps * eps_nats / ln2 * fine_per_bit) buckets – this bound. It composes additively with the existing steps-floor structural smear, so every count / rank / mass bracket in this package widens by exactly this many buckets and remains sound.

Practical reading: with the default quantizer (1/8-bit fine buckets) a per-step logit error of ~0.01 nats (typical of int8 inference) costs ceil(steps * 0.115) buckets – about one bucket per nine steps – while the forwards run 2-4x faster.

Parameters:

  • eps_nats (float)

  • steps (int)

  • quantizer (Quantizer)

Return type:

int

quantized_index(enum, max_bits, bin_width_bits=1.0)[source]

Convenience wrapper for QuantizedEnumerationIndex.from_enumerator.

Parameters:
Return type:

QuantizedEnumerationIndex

class QuantizedEnumerationIndex(bins, bin_width_bits, max_bits, truncated)[source]

Bases: object

Bounded, indexable view of an exact probability-ordered enumeration.

Items are grouped into log-probability bins measured in bits:

bits(x) = -log2(p(x)) bin(x) = floor(bits(x) / bin_width_bits)

Only items whose bit cost is at most max_bits are indexed. The returned (value, log_prob) pairs always carry the original exact log probability; the quantized bins are used only for counting and rank lookup.

The generic implementation can build from an exact enumerator stream, so it works for every currently enumerable discrete stats model. Simple distributions may also build directly from scored support items, and distribution-specific dynamic programs can produce the same bin layout without relying on an exact global stream.

Parameters:
static bin_for_log_prob(log_prob, bin_width_bits=1.0)[source]

Return the quantized bit bin for a log probability.

Parameters:
Return type:

int

classmethod from_enumerator(enum, max_bits, bin_width_bits=1.0)[source]

Build an index from an exact non-increasing-probability enumerator.

Parameters:

  • enum (Iterator[tuple[Any, float]]) – Iterator yielding (value, log_prob) pairs in exact descending order.

  • max_bits (float) – Include only values with -log2(p) <= max_bits.

  • bin_width_bits (float) – Width of each quantized probability bin in bits.

Returns:

QuantizedEnumerationIndex over the bounded prefix/domain slice.

Return type:

QuantizedEnumerationIndex

classmethod from_items(items, max_bits, bin_width_bits=1.0, sorted_items=False, truncated=None)[source]

Build an index from known support items and exact log probabilities.

Parameters:

  • items (Sequence[tuple[Any, float]]) – Finite collection of (value, log_prob) pairs. Zero-probability items with log_prob == -inf are ignored.

  • max_bits (float) – Include only values with -log2(p) <= max_bits.

  • bin_width_bits (float) – Width of each quantized probability bin in bits.

  • sorted_items (bool) – If True, preserve the input order. Otherwise, stable-sort included items by descending log probability.

  • truncated (bool | None) – Optional explicit truncation flag. If omitted, it is True when any finite-probability item was outside the bit bound.

Returns:

QuantizedEnumerationIndex over the bounded finite item set.

Return type:

QuantizedEnumerationIndex

bin_for_index(index)[source]

Return (bin_id, offset_within_bin) for a bounded quantized rank.

Parameters:

index (int)

Return type:

tuple[int, int]

get(index)[source]

Return the indexed (value, exact_log_prob) pair.

Parameters:

index (int)

Return type:

tuple[Any, float]

slice(start, k)[source]

Return up to k indexed pairs starting at start.

Parameters:
Return type:

list[tuple[Any, float]]

iter_from(start=0)[source]

Iterate indexed pairs from start to the end of the bounded index.

Parameters:

start (int)

Return type:

Iterator[tuple[Any, float]]

bin_items(bin_id)[source]

Return all indexed items in a quantized probability bin.

Parameters:

bin_id (int)

Return type:

list[tuple[Any, float]]

summary()[source]

Return a compact description of the bounded index.

Return type:

dict[str, Any]

class LazyQuantizedEnumerationIndex(counts, bin_width_bits, max_bits, truncated, getter)[source]

Bases: QuantizedEnumerationIndex

Quantized index whose bin counts are precomputed but items are unranked lazily.

This supports compositional distributions where a dynamic program can count how many values fall in each quantized log-density bin without materializing every Cartesian-product value. The getter receives a bin id and offset within that bin and returns the exact (value, log_prob) pair for that quantized rank.

Parameters:
bin_for_index(index)[source]

Return (bin_id, offset_within_bin) for a bounded quantized rank.

Parameters:

index (int)

Return type:

tuple[int, int]

get(index)[source]

Return the indexed (value, exact_log_prob) pair.

Parameters:

index (int)

Return type:

tuple[Any, float]

iter_from(start=0)[source]

Iterate indexed pairs from start to the end of the bounded index.

Parameters:

start (int)

Return type:

Iterator[tuple[Any, float]]

bin_items(bin_id)[source]

Return all indexed items in a quantized probability bin.

Parameters:

bin_id (int)

Return type:

list[tuple[Any, float]]

class CountSemiring[source]

Bases: DecomposableSemiring

Witness-retaining carrier over reified nodes / CountIndex: structural counting + unranking.

leaf/plus/scale/map_values build _CNode trees unranked by the iterative _unrank() (bounded call-stack regardless of chain depth); times/product convolve via mixle.enumeration.quantization.convolve_indices() (flat unranker – identical bin counts and within-bucket order to the previous path). Both element kinds expose .hist and .get_in_bucket, so they interoperate freely.

zero()[source]
Return type:

_CNode

one()[source]
Return type:

_CNode

leaf(value, log_prob, quantizer)[source]
Parameters:

  • value (Any)

  • log_prob (float)

  • quantizer (Quantizer)

Return type:

_CNode

from_enumerator(enum, quantizer, max_fine_bucket)[source]

Lift an atomic distribution’s exact enumerator into a carrier element (depth-bounded).

Parameters:
Return type:

tuple[CountIndex, bool]

plus(a, b)[source]
Return type:

_CNode

times(a, b, quantizer, max_fine_bucket)[source]
Parameters:

  • quantizer (Quantizer)

  • max_fine_bucket (int)

Return type:

CountIndex

product(elements, quantizer, max_fine_bucket)[source]

n-ary times. Default folds times; carriers may override for efficiency/order.

Parameters:

  • elements (Sequence[Any])

  • quantizer (Quantizer)

  • max_fine_bucket (int)

Return type:

CountIndex

scale(a, log_prob, quantizer, max_fine_bucket=None)[source]

Multiply by a constant probability factor: shift buckets by the factor, value unchanged.

This is the action of the base log-prob monoid on the carrier (e.g. a sequence length term, a Markov transition/initial term). Optionally truncate the shifted histogram to a depth bound.

Parameters:

  • log_prob (float)

  • quantizer (Quantizer)

  • max_fine_bucket (int | None)

Return type:

_CNode

map_values(a, fn)[source]

Relabel values (pushforward) without touching counts or buckets – e.g. tuple -> list.

Parameters:

fn (Callable[[Any], Any])

Return type:

_CNode

power_prefix(a, max_k, quantizer, max_fine_bucket)[source]

Return [a^(times 0), …, a^(times K)] (K <= max_k), the k-fold self-products.

The histograms are built incrementally (O(max_k) convolutions, shared), so this is the count side of an iid sequence. Each element’s unranker is the flat product of k copies, built lazily and cached – identical bin counts and within-bucket order to calling product([a]*k) directly, but without the O(k^2) eager cost. Stops early if the element mass is exhausted within the depth bound.

Parameters:

  • a (CountIndex)

  • max_k (int)

  • quantizer (Quantizer)

  • max_fine_bucket (int)

Return type:

Sequence[CountIndex]

class DecomposableSemiring[source]

Bases: ABC

A witness-retaining semiring over which a likelihood reduction is evaluated.

zero/one are the additive/multiplicative identities; plus pools mutually exclusive alternatives (e.g. different sequence lengths or Markov end-states); times composes independent factors (whose log-probabilities add); product is n-ary times. leaf lifts a single scored atom into the carrier. Implementations choose the carrier E; for counting it is CountIndex, whose per-bucket structure is the witness needed to unrank.

abstractmethod zero()[source]
Return type:

E

abstractmethod one()[source]
Return type:

E

abstractmethod leaf(value, log_prob, quantizer)[source]
Parameters:

  • value (Any)

  • log_prob (float)

  • quantizer (Quantizer)

Return type:

E

abstractmethod plus(a, b)[source]
Parameters:

  • a (E)

  • b (E)

Return type:

E

abstractmethod times(a, b, quantizer, max_fine_bucket)[source]
Parameters:

  • a (E)

  • b (E)

  • quantizer (Quantizer)

  • max_fine_bucket (int)

Return type:

E

product(elements, quantizer, max_fine_bucket)[source]

n-ary times. Default folds times; carriers may override for efficiency/order.

Parameters:

  • elements (Sequence[E])

  • quantizer (Quantizer)

  • max_fine_bucket (int)

Return type:

E

class TropicalSemiring[source]

Bases: DecomposableSemiring

Max-plus (Viterbi) carrier: the single best configuration and its log-probability.

The second realization of DecomposableSemiring – it swaps only the carrier (the contract this module was built around). Where CountSemiring counts the support in each quantized bin, the tropical carrier folds the same leaf/plus/times reduction over the (max, +) semiring: plus keeps the higher-probability alternative, times/product add log-probabilities and concatenate witnesses. The result is the Viterbi (most-probable) configuration and its exact log-probability – the top of the descending-probability order, and hence an exact bound usable by any family expressible in leaf/plus/times.

Crucially this includes the non-decomposable families (HMM / Mixture): exact counting couples across structure (which is why the count index does not serve them), but the Viterbi bound factors cleanly through the tropical fold. (Exact bounded counting for those families is the state-augmented count DP – a further step; for HMM the exact descending enumeration is already available via mixle.enumeration.hmm_paths.hmm_best_paths().)

The witness is the flat tuple of leaf atoms along the best configuration (the tropical carrier keeps a single witness, not the per-child structured value the count carrier reconstructs on unrank), which is exactly what a Viterbi decode wants. quantizer / max_fine_bucket are accepted for interface symmetry but unused – the tropical fold needs no quantization.

zero()[source]
Return type:

_Trop

one()[source]
Return type:

_Trop

leaf(value, log_prob, quantizer=None)[source]
Parameters:

  • value (Any)

  • log_prob (float)

  • quantizer (Quantizer | None)

Return type:

_Trop

plus(a, b)[source]
Parameters:

  • a (_Trop)

  • b (_Trop)

Return type:

_Trop

times(a, b, quantizer=None, max_fine_bucket=None)[source]
Parameters:

  • a (_Trop)

  • b (_Trop)

  • quantizer (Quantizer | None)

  • max_fine_bucket (int | None)

Return type:

_Trop

product(elements, quantizer=None, max_fine_bucket=None)[source]

n-ary times. Default folds times; carriers may override for efficiency/order.

Parameters:

  • elements (Sequence[_Trop])

  • quantizer (Quantizer | None)

  • max_fine_bucket (int | None)

Return type:

_Trop

best_first_union(streams, log_offsets, exact_log_density, tol=1.0e-10)[source]

Enumerate the union of sorted streams with overlapping supports.

Candidate values are pulled from the streams (stream k shifted by log_offsets[k]), de-duplicated via freeze, re-scored exactly with exact_log_density, and buffered until their exact score is at least the upper bound on any not-yet-seen value: logsumexp_k(log_offsets[k] + head_lp_k). This is the mixture algorithm: any unseen x satisfies p_k(x) <= head_k for every k, hence sum_k w_k p_k(x) <= exp(bound).

Parameters:
Return type:

Iterator[tuple[Any, float]]

merge_enumerators(streams, offsets)[source]

Lazy k-way merge of sorted (value, log_prob) streams with per-stream offsets.

Stream k’s log probs are shifted by offsets[k]. Correct only when the streams have pairwise disjoint supports (no de-duplication or re-scoring is performed).

Parameters:
Return type:

Iterator[tuple[Any, float]]

class ProductEnumerator(streams, combine=tuple, offset=0.0)[source]

Bases: object

Best-first enumeration of the Cartesian product of sorted child streams.

Yields (combine(values), log_prob) with log_prob = offset + sum of child log probs, in non-increasing order. Standard k-best lattice search: a max-heap over index tuples, with successors advancing one coordinate; correctness follows from each child stream being sorted (coordinate-wise monotonicity).

Duplicate-free successor rule: each heap entry carries the lowest coordinate it is allowed to advance, and a node only spawns successors for coordinates at or above that index. Every tuple is then reached by exactly one path – the one that increments coordinates in non-decreasing index order – so no visited set (and no O(n) per-successor tuple hashing) is needed. Best-first order is preserved because each successor’s score is <= its generating node’s (a sorted coordinate only decreases when advanced), so canonical edges are score-monotone.

Parameters:
sound_top_k(dist, k, start=0, budget_bits=40.0, oversample=8, bin_width_bits=1.0, max_budget_bits=512.0, total_mass=1.0, tol=1.0e-12)[source]

Exact true-descending observations ranked [start, start+k) for ANY normalized model.

Mass-threshold certificate, correct regardless of the seek stream’s ordering (so it holds for deep/nested mixtures and HMMs where the tropical order is badly displaced): pull distinct (value, log_prob) from the count-budget index, accumulate exact probability mass, and keep the best start+k by probability. Since every unpulled item’s probability is at most the remaining mass total_mass - accumulated, once remaining drops below the (start+k)-th best probability the heap is exactly the true top start+k; return sorted_desc[start:start+k]. If the in-budget stream is exhausted first, the budget is doubled (up to max_budget_bits). total_mass (default 1.0) may be a known upper bound for sub-normalized models.

Cost is O(number pulled) – small for peaked distributions, large for flat ones, so for top-k from rank 0 dist.enumerator() is usually leaner; this adds the soundness certificate and the arbitrary start offset the sequential enumerator lacks.

Parameters:
Return type:

list[tuple[Any, float]]

quantized_best_first_decode(next_logprobs=None, eos=None, max_len=None, top_k=None, top_p=None, bucket_bits=12, batch_next_logprobs=None, batch_size=64, start=(), max_results=None, min_mass=None)[source]

Fast descending-probability sequence enumeration specialized for neural / transformer decoders.

Three structure-aware accelerations over best_first_decode():

  1. Nucleus / top-k pruning. Neural next-token distributions are sharply peaked, so each step is restricted to its top_k tokens or its top_p nucleus – dropping the long low-probability tail collapses the branching factor (a ~50k vocab down to a handful) at negligible mass loss.

  2. Quantized bucket priority queue. Cumulative log-probs only decrease, so instead of an O(log n) comparison heap the frontier is bucketed by quantized score (bucket = floor(score * 2**bucket_bits)) and drained highest-bucket first – O(1) pushes/pops, and prefixes of near-equal score are grouped. Buckets are disjoint score ranges, so order is exact across buckets and within ~2**-bucket_bits inside one.

  3. Batched scoring. The cost is dominated by model forward passes. Pass batch_next_logprobs to score up to batch_size frontier prefixes in one call (one padded GPU forward) instead of one at a time.

With top_k=top_p=None and a large bucket_bits this reduces to the exact enumeration; pruning is the only approximation (report min_mass to stop once enough probability is covered).

Parameters:

  • next_logprobs (Callable[[tuple], Iterable[tuple[Any, float]]] | None) – next_logprobs(prefix) -> [(token, log_prob), ...] (used when batch_next_logprobs is not given). Log-probs must be <= 0.

  • eos (Any) – end-of-sequence token (a prefix ending in eos is complete).

  • max_len (int | None) – maximum sequence length. Give eos and/or max_len.

  • top_k (int | None) – keep only the top_k highest-probability tokens per step.

  • top_p (float | None) – keep the smallest set of tokens per step whose probability sums to >= top_p (nucleus).

  • bucket_bits (int) – score-quantization resolution; larger = finer ordering, slower bookkeeping.

  • batch_next_logprobs (Callable[[list[tuple]], list[Iterable[tuple[Any, float]]]] | None) – optional batch_next_logprobs([prefix, ...]) -> [[(token, log_prob), ...], ...] scoring a batch of prefixes in one forward pass.

  • batch_size (int) – number of frontier prefixes expanded per (batched) scoring call.

  • start (tuple) – initial prefix.

  • max_results (int | None) – stop after this many complete sequences.

  • min_mass (float | None) – stop once the yielded sequences cover at least this much probability mass.

Yields:

(sequence_tuple, total_log_prob), highest probability first (exact across score buckets).

Return type:

Iterator[tuple[tuple, float]]

best_first_decode(next_logprobs, eos=None, max_len=None, start=(), heuristic=None, max_results=None)[source]

Exactly enumerate an autoregressive model’s sequences in descending total log-probability.

Parameters:

  • next_logprobs (Callable[[tuple], Iterable[tuple[Any, float]]]) – next_logprobs(prefix) returns an iterable of (token, log_prob) continuations of prefix (e.g. the log-softmax of a transformer’s next-token logits). Log-probs must be <= 0.

  • eos (Any) – end-of-sequence token; a prefix whose last token is eos is complete (and not extended).

  • max_len (int | None) – maximum sequence length; a prefix of this length is complete. At least one of eos / max_len should be given or enumeration may not terminate.

  • start (tuple) – the initial prefix (default empty).

  • heuristic (Callable[[tuple], float] | None) – optional admissible upper bound on the remaining log-probability from a prefix (e.g. remaining_steps * max_step_logprob); tightens the search. Omit for the exact h=0 search.

  • max_results (int | None) – stop after this many complete sequences.

Yields:

(sequence_tuple, total_log_prob) in nonincreasing total log-probability.

Return type:

Iterator[tuple[tuple, float]]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]
Parameters:
Return type:

int

sampler(seed=None)[source]
Parameters:

seed (int | None)

Return type:

_ARSampler

enumerator()[source]

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

seek_index(*, max_depth_bits=4096.0)[source]

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)[source]

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

Parameters:
envelope_index(*, n_paths=64, seed=0, budget_bits=64.0)[source]

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:
top_k(k)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

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)[source]

CountDPSeekResult at descending index (with a bracket).

Parameters:

index (int)

rank(sequence)[source]

DensityRankResult – rank + cumulative mass of a sequence.

Parameters:

sequence (Iterable[Any])

cumulative(sequence)[source]

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

Parameters:

sequence (Iterable[Any])

nucleus_size(p)[source]

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

Parameters:

p (float)

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

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:
Return type:

tuple[CountIndex, bool]

class AREnvelopeIndex(model, *, n_paths=64, seed=0, budget_bits=64.0, calibration_sequences=None)[source]

Bases: object

Envelope (mean-field) seek index over a fixed-length AutoregressiveEnumerable.

Parameters:

  • model (Any) – the autoregressive adapter (max_len set; terminating/eos models are rejected).

  • n_paths (int) – contexts sampled per depth for the envelope calibration (more = sharper envelope; the root depth is always exact). The calibration forwards are memoized on the model, so they are shared with any exact index built later.

  • seed (int) – calibration sampling seed (the index is deterministic given it).

  • budget_bits (float) – initial depth of the suffix tables; queries deepen geometrically as needed.

  • calibration_sequences (list[tuple] | None) – optional typical sequences (a corpus, a provider’s fast generations) to calibrate the envelope from INSTEAD of ancestral sampling. Each is harvested through model.harvest – with the all_position_logprobs contract that is ONE forward per sequence for all its per-depth contexts, ~L-times cheaper than sampling token by token.

ensure_bits(depth_bits)[source]

Deepen the suffix tables to cover depth_bits (geometric, cheap: L capped convolutions).

Parameters:

depth_bits (float)

Return type:

AREnvelopeIndex

total()[source]

Estimated number of sequences within the built depth (exact for an iid-step model).

Return type:

float

count(min_log_prob)[source]

Estimated number of sequences with log_density >= min_log_prob (mean-field; iid-exact).

Parameters:

min_log_prob (float)

Return type:

float

mass_above(min_log_prob)[source]

A mean-field (lower, upper) estimate of the head mass above min_log_prob.

Same bucket arithmetic as the exact index’s mass_above (each bucket of c sequences holds between c * 2**-hi_bits and c * 2**-lo_bits of mass), applied to the envelope histogram – so the bracket carries the envelope’s estimation error on top of the quantization smear.

Parameters:

min_log_prob (float)

Return type:

tuple[float, float]

unrank(i)[source]

The approximately-i-th most probable sequence and its exact log-probability.

Costs one model forward per step (L total, memoized) – never a tree expansion. The rank coordinate inherits the envelope approximation; the returned sequence and its log-probability are exact model quantities. Raises IndexError past the estimated support size.

Parameters:

i (int)

Return type:

tuple[tuple, float]

threshold(rank)[source]

Exact log-probability of the sequence the envelope places at rank (approximate boundary).

Parameters:

rank (int)

Return type:

float

rank_bracket(sequence)[source]

Estimated [lo, hi] rank bracket of sequence – its envelope bucket’s rank span.

lo counts the estimated sequences in strictly shallower buckets; hi adds the sequence’s own bucket. Exact for iid-step models; otherwise a mean-field estimate (floats, not certificates).

Parameters:

sequence (Any)

Return type:

tuple[float, float]

class LatticeEnvelopeIndex(model, *, n_clusters=None, cluster_fn=None, n_paths=128, seed=0, budget_bits=64.0)[source]

Bases: object

Cluster-conditioned envelope index: the Markov refinement between mean-field and exact.

AREnvelopeIndex averages one envelope per depth – exact only for iid steps. This index conditions each depth’s envelope on the cluster of the last token (cluster_fn): the calibration histograms are kept per (depth, cluster) and split by the next token’s cluster, and the suffix tables become the lattice DP S_d[c] = sum_c' H_d[c][c'] (*) S_{d+1}[c'] – the same shape as HMMPathIndex with clusters as states. The result is exact for any order-1 Markov model whose next-token distribution depends only on (depth, cluster_fn(last token)) – with cluster_fn the identity, that is every last-token Markov model, where the mean-field envelope is provably lossy. Coarser clusterings interpolate: m = 1 recovers mean-field, m = V conditions on the full last token.

Queries mirror AREnvelopeIndex: counts and thresholds off the root table, unrank descends the REAL model apportioning by cluster-conditioned subtree estimates (O(L) forwards), every returned log-probability exact. A (depth, cluster) never visited in calibration borrows the depth’s aggregate envelope (mean-field fallback for that pocket) – raise n_paths to shrink the pockets. Fixed-length models only.

Parameters:

  • model (Any)

  • n_clusters (int | None)

  • cluster_fn (Any)

  • n_paths (int)

  • seed (int)

  • budget_bits (float)

ensure_bits(depth_bits)[source]
Parameters:

depth_bits (float)

Return type:

LatticeEnvelopeIndex

total()[source]

Estimated sequences within the built depth (exact for depth+last-cluster Markov models).

Return type:

float

count(min_log_prob)[source]

Estimated number of sequences with log_density >= min_log_prob (Markov-refined estimate).

Parameters:

min_log_prob (float)

Return type:

float

unrank(i)[source]

The approximately-i-th most probable sequence with its exact log-probability.

Parameters:

i (int)

Return type:

tuple[tuple, float]

threshold(rank)[source]
Parameters:

rank (int)

Return type:

float

rank_bracket(sequence)[source]

Estimated [lo, hi] rank bracket – exact for depth+last-cluster Markov models.

Parameters:

sequence (Any)

Return type:

tuple[float, float]

class RescoredIndex(draft_index, target, *, rerank_window=64, assumed_gap=None)[source]

Bases: object

Draft-ordered, target-scored enumeration with window reranking.

Parameters:

  • draft_index (Any) – any index with unrank(i) -> (sequence, draft_log_prob) – a SeekIndex over a cheap AutoregressiveEnumerable, an AREnvelopeIndex, or anything equivalent.

  • target (Any) – the expensive model – an AutoregressiveEnumerable (its score_sequences() batch scorer is used) or a bare callable [seqs] -> log_probs.

  • rerank_window (int) – extra draft candidates pulled around a query and reranked by target score. Larger = more robust to draft/target disagreement, one batched forward either way.

  • assumed_gap (float | None) – optional global bound on |target_lp - draft_lp| (nats). When supplied, results carry a sound certified verdict; otherwise certified is None and the observed gap is reported as a diagnostic.

top_k(k)[source]

The k best sequences by TARGET score among the k + rerank_window draft head.

Returns {"items": [(seq, target_lp), ...], "certified": bool | None, "gap": float} – target-exact scores, draft+window-approximate completeness (see the class docstring).

Parameters:

k (int)

Return type:

dict[str, Any]

slice(start, k)[source]

Target-reranked [start, start + k) slice of the pulled start + k + rerank_window head.

Same semantics as top_k(): order within the pulled set is target-exact; the certificate covers whether an unpulled sequence could belong in (or before) the slice.

Parameters:
Return type:

dict[str, Any]

unrank(i)[source]

The draft’s rank-i sequence with the TARGET’s exact log-probability.

The rank coordinate is the draft’s (no reranking): the cheap random-access primitive. Use top_k() / slice() when local target-order matters.

Parameters:

i (int)

Return type:

tuple[tuple, float]

hmm_best_paths(log_pi, log_A, log_b, k=None)[source]

Enumerate HMM state paths in nonincreasing joint log-probability.

Parameters:

  • log_pi (ndarray) – (K,) log initial-state distribution.

  • log_A (ndarray) – (K, K) log transition matrix (row j -> column k is log p(z_t=k | z_{t-1}=j)).

  • log_b (ndarray) – (T, K) per-position emission log-likelihoods log p(x_t | z_t=k).

  • k (int | None) – stop after the k best paths; None enumerates all K**T lazily.

Yields:

(path, joint_log_prob) with path a length-T tuple of state indices, highest first. The first yield is the Viterbi (MAP) path.

Return type:

Iterator[tuple[tuple[int, …], float]]

class HMMPathIndex(log_pi, log_A, log_b, *, bin_width_bits=1.0, oversample=8, budget_bits=None)[source]

Bases: object

Quantized random-access index over an HMM’s state paths for one observation sequence.

Precompute (once, O(T * K^2 * W)): quantize every step score – log_pi[s] + log_b[0, s] and log_A[s, s'] + log_b[t, s'], one floor per step so the accumulated smear is at most T fine buckets – and run a forward count DP: C_t[s'] is the histogram, over integer total-score buckets, of the number of length-t+1 prefixes ending in state s'. Counts are float64 (exact below 2**53; the number of paths is K**T, so deep problems carry the documented ~1e-16/op relative error instead of overflowing).

Query (each O(T * K)): unrank(i) walks the stored tables backward – pick the final state whose bucket count covers the offset, then repeatedly the predecessor whose shifted prefix bucket does – returning the i-th best path by quantized score and its exact joint log-probability. count / threshold / mass_above read the pooled final histogram.

Ordering contract: paths are ordered by their quantized bucket (width bin_width_bits/oversample bits); within a bucket the order is deterministic but unspecified – exactly the count-index semantics elsewhere in mixle.enumeration. hmm_best_paths() remains the exact-order tool for the head; this index is the random-access tool for depth (rank 1e6 costs one table walk, not 1e6 A* expansions).

Parameters:

  • log_pi (np.ndarray)

  • log_A (np.ndarray)

  • log_b (np.ndarray)

  • bin_width_bits (float)

  • oversample (int)

  • budget_bits (float | None)

bucket_of(log_joint)[source]

The index’s bucket frame for a true joint score (bits behind the per-position optimum).

Parameters:

log_joint (float)

Return type:

int

total()[source]

Number of state paths within the budget (== K**T when nothing truncated; float64 counts).

Return type:

float

count(min_log_joint)[source]

How many paths have quantized joint log-probability at least min_log_joint.

Counts every true qualifier (the structural bucket never over-states a path’s bits) plus at most the paths within the T-floor smear band below the threshold.

Parameters:

min_log_joint (float)

Return type:

float

mass_above(min_log_joint)[source]

A (lower, upper) bracket on the total joint probability/density of paths above the threshold.

Bucket arithmetic with the T-floor smear: a path in bucket b carries between exp(total_offset) * 2**-((b + T) / fpb) and exp(total_offset) * 2**-(b / fpb) of joint mass (the offset restores the per-position shift, so unnormalized emission likelihoods work).

Parameters:

min_log_joint (float)

Return type:

tuple[float, float]

unrank(i)[source]

The i-th best state path by quantized score (0-based) and its exact joint log-probability.

One backward table walk – O(T * K) – regardless of how deep i is.

Parameters:

i (int)

Return type:

tuple[tuple[int, …], float]

threshold(rank)[source]

Exact joint log-probability of the rank-th best (quantized-order) path.

Parameters:

rank (int)

Return type:

float

iter_paths(start=0)[source]

Iterate paths from quantized rank start (sequential unranks over the stored tables).

Parameters:

start (int)

Return type:

Iterator[tuple[tuple[int, …], float]]

Subpackages

Submodules