mixle.stats.latent package

Public latent-variable probability-model namespace.

The focused namespace exposes fitted families, estimators, samplers, enumerators, encoders, accumulators, operators, schedules, diagnostics, and high-level helpers. Exports are lazy so importing mixle.stats.latent does not initialize every optional numerical backend. The curated non-plumbing subset is also available from mixle.stats; integration roles remain directly importable from either package.

class LatentAPIEntry(name, source, status, lazy=True)[source]

Bases: object

One immutable entry in the reviewed latent-package export manifest.

Parameters:
latent_api_manifest()[source]

Return the immutable, ordered latent-package export manifest.

Return type:

tuple[LatentAPIEntry, …]

class ChainedAttentionDistribution(keys, emission, sigma2=0.1, name=None)[source]

Bases: SequenceEncodableProbabilityDistribution

An L-hop stack of responsibility-attention heads (chained, content-addressed).

Parameters:

  • keys (np.ndarray)

  • emission (np.ndarray)

  • sigma2 (float)

  • name (str | None)

density(x)[source]

Return the probability of one chained-attention observation.

Parameters:

x (tuple[Any, Any, int, int])

Return type:

float

log_density(x)[source]

Return the log-probability of one context/query/target observation.

Parameters:

x (tuple[Any, Any, int, int])

Return type:

float

seq_log_density(x)[source]

Return vectorized log-probabilities for encoded chained-attention observations.

Return type:

ndarray

predict_proba(context_keys, context_values, query)[source]

Predictive target distribution (target marginalized); (T,) or (n, T).

Parameters:
Return type:

ndarray

sampler(seed=None)[source]

Return a sampler for synthetic chained-attention observations.

Parameters:

seed (int | None)

Return type:

ChainedAttentionSampler

estimator(pseudo_count=None)[source]

Return a closed-form EM estimator for this attention chain.

Parameters:

pseudo_count (float | None)

Return type:

ChainedAttentionEstimator

dist_to_encoder()[source]

Return the encoder for context keys, values, query symbols, and targets.

Return type:

ChainedAttentionDataEncoder

class ChainedAttentionSampler(dist, seed=None)[source]

Bases: DistributionSampler

Generative sampler (uniform context + the chained gate).

Parameters:

  • dist (ChainedAttentionDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw one observation or size iid synthetic observations.

Parameters:
Return type:

Any

class ChainedAttentionAccumulator(n_hops, num_symbols, num_targets, keys=None, name=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Forward-backward sufficient statistics (additive): per-hop key numerators + emission counts.

seq_update(x, weights, estimate)[source]

Update forward-backward sufficient statistics from encoded observations.

Parameters:

estimate (ChainedAttentionDistribution)

Return type:

None

seq_initialize(x, weights, rng)[source]

Initialize sufficient statistics with random hop responsibilities.

Parameters:

rng (RandomState)

Return type:

None

update(x, weight, estimate)[source]

Update from one weighted chained-attention observation.

Parameters:

weight (float)

Return type:

None

initialize(x, weight, rng)[source]

Initialize from one weighted chained-attention observation.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge key-table, emission, likelihood, and weight statistics.

Return type:

ChainedAttentionAccumulator

value()[source]

Return key statistics, emission counts, log-likelihood, and total weight.

from_value(x)[source]

Restore accumulator state from value output.

Return type:

ChainedAttentionAccumulator

acc_to_encoder()[source]

Return the encoder compatible with this accumulator.

Return type:

ChainedAttentionDataEncoder

class ChainedAttentionAccumulatorFactory(n_hops, num_symbols, num_targets, keys=None, name=None)[source]

Bases: StatisticAccumulatorFactory

Create accumulators for chained-attention EM statistics.

make()[source]

Create an empty chained-attention accumulator.

Return type:

ChainedAttentionAccumulator

class ChainedAttentionEstimator(n_hops, num_symbols, num_targets, *, sigma2=0.1, emission_smoothing=1e-4, pseudo_count=None, name=None, keys=None)[source]

Bases: ParameterEstimator

Closed-form EM estimator: per-hop key tables (GMM means of one-hot queries) + emission counts.

Parameters:

  • n_hops (int)

  • num_symbols (int)

  • num_targets (int)

  • sigma2 (float)

  • emission_smoothing (float)

  • pseudo_count (float | None)

  • name (str | None)

  • keys (str | None)

accumulator_factory()[source]

Return a factory for chained-attention sufficient-statistic accumulators.

Return type:

ChainedAttentionAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate key tables and emissions from accumulated forward-backward statistics.

Parameters:

nobs (float | None)

Return type:

ChainedAttentionDistribution

class ChainedAttentionDataEncoder(num_symbols=None, num_targets=None)[source]

Bases: DataSequenceEncoder

Encodes (context_keys, context_values, query_symbol, target) into stacked integer arrays.

seq_encode(x)[source]

Encode (context_keys, context_values, query, target) observations.

Parameters:

x (Sequence[tuple[Any, Any, int, int]])

Return type:

tuple[ndarray, ndarray, ndarray, ndarray]

row_count(x)[source]

Return the number of observations represented by x.

The default covers simple array/list encodings and structurally aligned tuples/dicts. Encoders with flattened, ragged, or otherwise specialized layouts must override this method so metadata wrappers can verify counts instead of trusting a caller assertion.

Return type:

int

class DiracLengthMixtureDistribution(len_dist, p, v=0, name=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Mixture between a fixed length and a learned length distribution.

Parameters:

  • p (float) – Probability of being drawn from length distribution. Must be between 0 and 1.

  • len_dist (SequenceEncodableProbabilityDistribution) – Distribution with support on non-negative integers.

  • name (Optional[str]) – Optional distribution name.

  • v (int)

p

Probability of being drawn from length distribution. Must be between 0 and 1.

Type:

float

len_dist

Distribution with support on non-negative integers.

Type:

SequenceEncodableProbabilityDistribution

name

Optional distribution name.

Type:

Optional[str]

compute_capabilities()[source]

Declare generated-compute support inherited from the length distribution.

compute_declaration()[source]

Return the generated-compute declaration for the Dirac-length mixture.

density(x)[source]

Evaluate density of length Dirac mixture distribution at observation x.

See log_density() for details.

Parameters:

x (int) – Integer value.

Returns:

Density at x.

Return type:

float

log_density(x)[source]

Evaluate the log-density of length Dirac mixture distribution at observation x.

log(P(x)) = log( p*P_1(x) + (1-p)*Delta_{v}(x) ),

Parameters:

x (int) – Integer value.

Returns:

log-density at x.

Return type:

float

component_log_density(x)[source]

Log-density of each mixture component (length distribution, dirac at v) at x.

Parameters:

x (int) – Integer value.

Returns:

Numpy array of the two component log-densities.

Return type:

ndarray

posterior(x)[source]

Posterior probability of each mixture component given observation x.

Parameters:

x (int) – Integer value.

Returns:

Numpy array of the two component posterior probabilities (sums to one).

Return type:

ndarray

seq_component_log_density(x)[source]

Vectorized component log-densities at sequence encoded input x.

Parameters:

x (E) – Sequence encoded data from DiracLengthMixtureDataEncoder.

Returns:

Numpy array of shape (len(x), 2) of component log-densities.

Return type:

ndarray

seq_log_density(x)[source]

Vectorized evaluation of the mixture log-density at sequence encoded input x.

Parameters:

x (E) – Sequence encoded data from DiracLengthMixtureDataEncoder.

Returns:

Numpy array of log-density (float) of len(x).

Return type:

ndarray

backend_seq_component_log_density(x, engine)[source]

Engine-neutral component log densities for encoded length/dirac mixtures.

Parameters:
Return type:

Any

backend_seq_log_density(x, engine)[source]

Engine-neutral mixture log-density for encoded length/dirac observations.

Parameters:
Return type:

Any

classmethod backend_stacked_params(dists, engine)[source]

Return stacked parameters for shared-dirac length mixtures.

Parameters:

  • dists (Sequence[DiracLengthMixtureDistribution])

  • engine (Any)

Return type:

dict[str, Any]

classmethod backend_stacked_log_density(x, params, engine)[source]

Return an (n, k) matrix of length/dirac mixture log densities.

Parameters:
Return type:

Any

classmethod backend_stacked_sufficient_statistics_with_estimator(x, weights, params, engine, estimator)[source]

Return per-component legacy (component_counts, length_stat) statistics.

Parameters:
Return type:

tuple[Any, …]

seq_posterior(x)[source]

Vectorized component posterior probabilities at sequence encoded input x.

Parameters:

x (E) – Sequence encoded data from DiracLengthMixtureDataEncoder.

Returns:

Numpy array of shape (len(x), 2) of component posteriors.

Return type:

ndarray

sampler(seed=None)[source]

Create a DiracLengthMixtureSampler from parameters of this distribution.

Parameters:

seed (Optional[int]) – Used to set seed in random sampler.

Returns:

DiracLengthMixtureSampler object.

Return type:

DiracLengthMixtureSampler

estimator(pseudo_count=None)[source]

Create a DiracLengthMixtureEstimator with matching dirac value v.

Parameters:

pseudo_count (Optional[float]) – Used to inflate sufficient statistics.

Returns:

DiracLengthMixtureEstimator object.

Return type:

DiracLengthMixtureEstimator

dist_to_encoder()[source]

Returns a DiracLengthMixtureDataEncoder for encoding sequences of iid integer observations.

Return type:

DiracLengthMixtureDataEncoder

enumerator()[source]

Returns a DiracLengthMixtureEnumerator iterating the union of the length-distribution support and the dirac point v in descending probability order.

Return type:

DiracLengthMixtureEnumerator

class DiracLengthMixtureEnumerator(dist)[source]

Bases: DistributionEnumerator

Enumerates the union of the length-distribution support and the dirac point v.

The model is a two-component mixture: the length distribution with weight p and a dirac delta at v with weight 1-p. The dirac component contributes the trivial single-point stream [(v, 0.0)]. Supports may overlap (the length distribution can also emit v), so candidates are de-duplicated and re-scored exactly with the mixture log-density.

Parameters:

dist (DiracLengthMixtureDistribution)

class DiracLengthMixtureSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for a Dirac-length mixture distribution.

Parameters:

  • dist (DiracLengthMixtureDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw iid samples from a DiracLengthMixture distribution.

Parameters:

  • size (Optional[int]) – Number of iid samples to draw.

  • batched (bool)

Returns:

Int or List[int] depending on size = None or size (int).

Return type:

list[int] | int

class DiracLengthMixtureAccumulator(accumulator, v=0, keys=(None, None), name=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulate component counts and length-distribution statistics.

Parameters:

  • accumulator (SequenceEncodableStatisticAccumulator) – Accumulator for the length distribution.

  • v (int) – Dirac location.

  • keys (Tuple[Optional[str], Optional[str]]) – Keys for the mixture weights and component statistics.

  • name (Optional[str]) – Optional accumulator name.

accumulator

Accumulator for the length distribution.

Type:

SequenceEncodableStatisticAccumulator

comp_counts

Posterior-weighted counts for the two components.

Type:

np.ndarray

weight_key

Key for merging mixture weight counts.

Type:

Optional[str]

comp_key

Key for merging component sufficient statistics.

Type:

Optional[str]

v

Dirac location.

Type:

int

name

Optional accumulator name.

Type:

Optional[str]

seq_update(x, weights, estimate)[source]

Vectorized accumulation of posterior-weighted statistics from encoded observations x.

Parameters:

  • x (E) – Sequence encoded data from DiracLengthMixtureDataEncoder.

  • weights (np.ndarray) – Weights on the observations.

  • estimate (DiracLengthMixtureDistribution) – Previous estimate used for posteriors.

seq_update_engine(x, weights, estimate, engine)[source]

Engine-resident E-step: the length-distribution scoring (the heavy term) runs through the active engine; the low-overhead two-component (length vs. Dirac) responsibility bookkeeping mirrors the host seq_update exactly.

update(x, weight, estimate)[source]

Add one observation’s posterior-weighted contribution to the sufficient statistics.

Parameters:

  • x (int) – Integer observation.

  • weight (float) – Weight on the observation.

  • estimate (DiracLengthMixtureDistribution) – Previous estimate used for posteriors.

Return type:

None

initialize(x, weight, rng)[source]

Initialize the accumulator with observation x, randomly splitting weight at the dirac point.

Parameters:

  • x (int) – Integer observation.

  • weight (float) – Weight on the observation.

  • rng (RandomState) – Random number generator for initialization.

seq_initialize(x, weights, rng)[source]

Vectorized initialization from encoded observations x with random splits at the dirac point.

Parameters:

  • x (E) – Sequence encoded data from DiracLengthMixtureDataEncoder.

  • weights (np.ndarray) – Weights on the observations.

  • rng (RandomState) – Random number generator for initialization.

Return type:

None

combine(suff_stat)[source]

Combine sufficient statistics (component counts, length-dist stats) with this accumulator.

Parameters:

suff_stat (Tuple[np.ndarray, SS0]) – Component counts and length-distribution statistics.

Returns:

This DiracLengthMixtureAccumulator.

Return type:

DiracLengthMixtureAccumulator

value()[source]

Returns sufficient statistics as a tuple (component counts, length-distribution statistics).

Return type:

tuple[ndarray, Any]

from_value(x)[source]

Set sufficient statistics from a (component counts, length-distribution statistics) tuple.

Parameters:

x (Tuple[np.ndarray, SS0]) – Component counts and length-distribution statistics.

Returns:

This DiracLengthMixtureAccumulator.

Return type:

DiracLengthMixtureAccumulator

key_merge(stats_dict)[source]

Merge keyed sufficient statistics into stats_dict under the weight and component keys.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)[source]

Replace keyed sufficient statistics from stats_dict under the weight and component keys.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()[source]

Returns a DiracLengthMixtureDataEncoder for encoding sequences of iid integer observations.

Return type:

DiracLengthMixtureDataEncoder

class DiracLengthMixtureAccumulatorFactory(factory, v=0, keys=(None, None), name=None)[source]

Bases: StatisticAccumulatorFactory

Factory for Dirac-length mixture accumulators.

Parameters:

  • factory (StatisticAccumulatorFactory) – Accumulator factory for the length distribution.

  • v (int) – Dirac location.

  • keys (Tuple[Optional[str], Optional[str]]) – Keys for the mixture weights and component statistics.

  • name (Optional[str]) – Optional accumulator name.

factory

Accumulator factory for the length distribution.

Type:

StatisticAccumulatorFactory

v

Dirac location.

Type:

int

keys

Keys for the mixture weights and component statistics.

Type:

Tuple[Optional[str], Optional[str]]

name

Optional accumulator name.

Type:

Optional[str]

make()[source]

Returns a new DiracLengthMixtureAccumulator wrapping a fresh length-distribution accumulator.

Return type:

DiracLengthMixtureAccumulator

class DiracLengthMixtureEstimator(estimator, v=0, fixed_p=None, suff_stat=None, pseudo_count=None, name=None, keys=(None, None))[source]

Bases: ParameterEstimator

Estimate Dirac-length mixture distributions.

Parameters:

  • estimator (ParameterEstimator) – Estimator for the length distribution.

  • v (int) – Dirac location.

  • fixed_p (Optional[float]) – Hold the length-distribution weight p fixed at this value.

  • suff_stat (Optional[float]) – Prior value of p used with pseudo_count for regularization.

  • pseudo_count (Optional[float]) – Used to inflate the component count statistics.

  • name (Optional[str]) – Optional name assigned to estimated distributions.

  • keys (Tuple[Optional[str], Optional[str]]) – Keys for the mixture weights and component statistics.

estimator

Estimator for the length distribution.

Type:

ParameterEstimator

v

Dirac location.

Type:

int

pseudo_count

Used to inflate the component count statistics.

Type:

Optional[float]

suff_stat

Prior value of p used with pseudo_count for regularization.

Type:

Optional[float]

keys

Keys for the mixture weights and component statistics.

Type:

Tuple[Optional[str], Optional[str]]

name

Optional name assigned to estimated distributions.

Type:

Optional[str]

fixed_p_vec

Fixed component weights [p, 1-p] when fixed_p is given.

Type:

Optional[np.ndarray]

accumulator_factory()[source]

Returns a DiracLengthMixtureAccumulatorFactory consistent with this estimator.

Return type:

DiracLengthMixtureAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate a DiracLengthMixtureDistribution from accumulated sufficient statistics.

Parameters:

  • nobs (Optional[float]) – Weighted number of observations.

  • suff_stat (Tuple[np.ndarray, SS0]) – Component counts and length-distribution statistics.

Returns:

DiracLengthMixtureDistribution object.

Return type:

DiracLengthMixtureDistribution

class DiracLengthMixtureDataEncoder(encoder, v=0)[source]

Bases: DataSequenceEncoder

Data encoder for iid integer observations under a Dirac-length mixture.

Parameters:

  • encoder (DataSequenceEncoder) – Encoder for the length distribution.

  • v (int) – Dirac location.

encoder

Encoder for the length distribution.

Type:

DataSequenceEncoder

v

Dirac location.

Type:

int

row_count(x)[source]

Return the explicit observation count stored in the encoded payload.

Parameters:

x (tuple[int, ndarray, ndarray, Any])

Return type:

int

seq_encode(x)[source]

Encode a sequence of iid integer observations for vectorized use.

Parameters:

x (Sequence[int]) – Sequence of iid integer observations.

Returns:

Tuple of (sequence length, indices equal to v, indices not equal to v, base-encoded data).

Return type:

tuple[int, ndarray, ndarray, Any]

class EffectiveSampleReceipt(declared_mass, assigned_mass, unassigned_mass, allows_unassigned, schema_version=1, contract='mixle.latent_effective_sample/v1')[source]

Bases: object

Validated relationship between declared and assigned effective mass.

Parameters:

  • declared_mass (float | None)

  • assigned_mass (float)

  • unassigned_mass (float)

  • allows_unassigned (bool)

  • schema_version (int)

  • contract (str)

validated_positive_integer(value, label)[source]

Return a positive exact non-boolean integer.

Parameters:
Return type:

int

validated_observation_weight(value, label='observation weight')[source]

Return one finite non-negative, non-boolean observation weight.

Parameters:
Return type:

float

validated_observation_weights(values, rows, label='observation weights')[source]

Return an owned finite non-negative weight vector with exact row geometry.

Parameters:
Return type:

ndarray

validated_count_array(values, shape, label='sufficient-statistic counts')[source]

Return an owned finite non-negative count array with exact geometry.

Parameters:
Return type:

ndarray

validated_statistic_tuple(values, arity, label)[source]

Return a statistic tuple only when its public schema arity is exact.

Parameters:
Return type:

tuple[Any, …]

validate_effective_sample_mass(declared_mass, assigned_mass, *, label, allow_unassigned=False)[source]

Validate declared outer mass against mass assigned to latent children.

declared_mass=None means the caller did not provide outer-mass metadata. Otherwise exact agreement is required, except that families whose documented impossible-evidence policy assigns zero responsibility may set allow_unassigned=True. Assigned mass may never exceed declared mass.

Parameters:

  • declared_mass (Any)

  • assigned_mass (Any)

  • label (str)

  • allow_unassigned (bool)

Return type:

EffectiveSampleReceipt

validated_weighted_responsibilities(values, weights, components, *, label, allow_unassigned=False)[source]

Validate an owned (rows, components) matrix of assigned posterior mass.

Parameters:

  • values (Any)

  • weights (Any)

  • components (int)

  • label (str)

  • allow_unassigned (bool)

Return type:

ndarray

class GateOptimizationReceipt(converged, steps_completed, initial_loss, final_loss, gradient_norm, termination)[source]

Bases: object

Outcome of a bounded softmax-gate optimization.

Parameters:
class GateBufferReceipt(rows_seen, rows_retained, rows_dropped, capacity, selection='deterministic_hash_top_k')[source]

Bases: object

Bounded deterministic coreset status for gate-training rows.

Parameters:

  • rows_seen (int)

  • rows_retained (int)

  • rows_dropped (int)

  • capacity (int)

  • selection (str)

class GatedMixtureStatistics(component_statistics, covariates, responsibilities, buffer_receipt, component_counts, schema_version=1)[source]

Bases: Sequence[Any]

Versioned gated-mixture statistics with backward four-slot iteration.

Parameters:

  • component_statistics (tuple[Any, ...])

  • covariates (ndarray)

  • responsibilities (ndarray)

  • buffer_receipt (GateBufferReceipt)

  • component_counts (ndarray)

  • schema_version (int)

class SoftmaxGate(weight, bias, *, fit_receipt=None)[source]

Bases: object

A torch-free multinomial-logistic gate p(k | z) = softmax(W z + b)_k, fit on soft targets.

fit(Z, R) minimizes the soft cross-entropy -sum_{n,k} R_{n,k} log p(k | z_n) by gradient descent – R are the responsibilities (rows need not sum to 1; the sample weight is folded in).

Parameters:

  • weight (np.ndarray)

  • bias (np.ndarray)

  • fit_receipt (GateOptimizationReceipt | None)

classmethod zeros(n_classes, n_features)[source]

Create a zero-logit gate with uniform initial class probabilities.

Parameters:

  • n_classes (int)

  • n_features (int)

Return type:

SoftmaxGate

log_prob_batch(z)[source]

(n, K) log-gate log p(k | z_n) for each row of z (shape (n, p)).

Parameters:

z (ndarray)

Return type:

ndarray

fit_with_receipt(z, resp, *, steps=200, lr=0.1, tol=1.0e-8)[source]

Fit a softmax gate and return explicit convergence/failure status.

Parameters:
Return type:

tuple[SoftmaxGate, GateOptimizationReceipt]

fit(z, resp, *, steps=200, lr=0.1, tol=1.0e-8)[source]

Fit a softmax gate; the returned gate carries fit_receipt.

Parameters:
Return type:

SoftmaxGate

class GatedMixtureDistribution(components, gate, name=None, keys=None, *, gate_fit_receipt=None, gate_buffer_receipt=None)[source]

Bases: SequenceEncodableProbabilityDistribution

A mixture whose weights are a gate p(k | z); observations are (z, y) pairs.

Parameters:

  • components (Sequence[SequenceEncodableProbabilityDistribution])

  • gate (Any)

  • name (str | None)

  • keys (str | None)

  • gate_fit_receipt (GateOptimizationReceipt | None)

  • gate_buffer_receipt (GateBufferReceipt | None)

log_density(x)[source]

Return log p(y | z) for one covariate/response pair.

Parameters:

x (tuple[Any, Any])

Return type:

float

seq_log_density(enc)[source]

Return vectorized conditional log-densities for encoded (z, y) pairs.

Parameters:

enc (Any)

Return type:

ndarray

posterior(x)[source]

Return posterior expert responsibilities for one (z, y) observation.

Parameters:

x (tuple[Any, Any])

Return type:

ndarray

sampler(seed=None)[source]

Return a conditional sampler that requires a covariate z.

Parameters:

seed (int | None)

Return type:

GatedMixtureSampler

estimator(pseudo_count=None)[source]

Return an EM estimator for experts and the covariate-dependent gate.

Parameters:

pseudo_count (float | None)

Return type:

GatedMixtureEstimator

dist_to_encoder()[source]

Return the encoder for covariates plus expert response encodings.

Return type:

GatedMixtureDataEncoder

class GatedMixtureSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sample y given a supplied z: draw a component from gate(z), then sample that expert.

Parameters:

  • dist (GatedMixtureDistribution)

  • seed (int | None)

sample_given(z)[source]

Sample a response from the gated mixture conditional on covariate z.

Parameters:

z (Any)

Return type:

Any

sample(size=None, *, batched=True)[source]

Raise because unconditional sampling requires caller-supplied covariates.

Parameters:
Return type:

Any

class GatedMixtureDataEncoder(component_encoders, n_features=None)[source]

Bases: DataSequenceEncoder

Encode [(z, y), ...] as (z array (n, p), per-expert encodings of the y column).

Parameters:

  • component_encoders (Sequence[DataSequenceEncoder])

  • n_features (int | None)

seq_encode(data)[source]

Encode covariates as a dense matrix and responses for every expert.

Parameters:

data (Sequence[tuple[Any, Any]])

Return type:

tuple[ndarray, tuple[Any, …]]

class GatedMixtureAccumulator(component_accumulators, num_components, keys=None, *, max_buffer_rows=10_000, n_features=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Route expert statistics and retain a bounded mergeable gate-training coreset.

Parameters:

  • component_accumulators (Sequence[Any])

  • num_components (int)

  • keys (str | None)

  • max_buffer_rows (int)

  • n_features (int | None)

property buffer_receipt: GateBufferReceipt

Return current bounded-buffer retention status.

seq_update(enc, weights, estimate)[source]

Update expert accumulators and gate buffers from encoded observations.

Parameters:

  • enc (Any)

  • weights (ndarray)

  • estimate (GatedMixtureDistribution | None)

Return type:

None

seq_initialize(enc, weights, rng)[source]

Initialize expert accumulators with random responsibility allocations.

Parameters:
Return type:

None

update(x, weight, estimate)[source]

Update from one weighted (z, y) observation.

Parameters:
Return type:

None

initialize(x, weight, rng)[source]

Initialize from one weighted (z, y) observation.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge expert sufficient statistics and buffered gate training data.

Parameters:

suff_stat (tuple[Any, ...])

Return type:

GatedMixtureAccumulator

value()[source]

Return owned expert statistics, bounded gate rows, and retention status.

Return type:

GatedMixtureStatistics

from_value(x)[source]

Restore expert statistics and gate training buffers.

Parameters:

x (tuple[Any, ...])

Return type:

GatedMixtureAccumulator

key_merge(stats_dict)[source]

Delegate keyed merges to expert accumulators.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)[source]

Delegate keyed replacements to expert accumulators.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()[source]

Return the encoder composed from expert accumulator encoders.

Return type:

GatedMixtureDataEncoder

class GatedMixtureAccumulatorFactory(component_factories, num_components, keys=None, *, max_buffer_rows=10_000, n_features=None)[source]

Bases: StatisticAccumulatorFactory

Create accumulators for gated-mixture EM.

Parameters:

  • component_factories (Sequence[Any])

  • num_components (int)

  • keys (str | None)

  • max_buffer_rows (int)

  • n_features (int | None)

make()[source]

Create an empty gated-mixture accumulator.

Return type:

GatedMixtureAccumulator

class GatedMixtureEstimator(component_estimators, gate, gate_steps=200, gate_lr=0.1, gate_tol=1.0e-8, max_buffer_rows=10_000, name=None, keys=None)[source]

Bases: ParameterEstimator

M-step: refit each expert from its responsibility-weighted stats, refit the gate on (z, resp).

Parameters:

  • component_estimators (Sequence[ParameterEstimator])

  • gate (Any)

  • gate_steps (int)

  • gate_lr (float)

  • gate_tol (float)

  • max_buffer_rows (int)

  • name (str | None)

  • keys (str | None)

accumulator_factory()[source]

Return a factory for gated-mixture sufficient-statistic accumulators.

Return type:

GatedMixtureAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate experts from responsibility-weighted stats and refit the gate.

Parameters:
Return type:

GatedMixtureDistribution

class GaussianMixtureDistribution(mu, sig2, w=MISSING, name=None, weights=MISSING)[source]

Bases: MixtureDistribution

Finite mixture of full-covariance multivariate Gaussian components.

The constructor packs each (mu[k], sig2[k]) pair into a MultivariateGaussianDistribution. Density, posterior responsibility, and vectorized scoring behavior are inherited from MixtureDistribution.

Parameters:
mu

Component means as a (K, d) NumPy array.

sig2

Component covariance matrices as a (K, d, d) NumPy array.

dim

Observation dimension.

density_semantics()[source]

Return exact-or-approximate density semantics joined from Gaussian components.

sampler(seed=None)[source]

Return a sampler that draws vectors from this Gaussian mixture.

Parameters:

seed (int | None) – Optional seed for reproducible component and Gaussian draws.

Returns:

GaussianMixtureSampler bound to this distribution.

Return type:

GaussianMixtureSampler

estimator(pseudo_count=None)[source]

Return an estimator with matching Gaussian component structure.

Parameters:

pseudo_count (float | None) – Optional smoothing mass for mixture weights.

Returns:

GaussianMixtureEstimator suitable for fitting length-d vectors.

Return type:

GaussianMixtureEstimator

dist_to_encoder()[source]

Return the Gaussian mixture encoder for iid vector observations.

Return type:

GaussianMixtureDataEncoder

class GaussianMixtureSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for iid draws from a Gaussian mixture.

Parameters:

  • dist (GaussianMixtureDistribution) – Gaussian mixture distribution to sample from.

  • seed (int | None) – Optional seed for reproducible component and Gaussian draws.

rng

Random state used for component labels.

compSamplers

Per-component Gaussian samplers.

sample(size=None, *, batched=True)[source]

Draw iid samples from the Gaussian mixture.

If size is None, a single length-d numpy array is returned. Otherwise a numpy array with shape (size, d) containing ‘size’ iid samples is returned. With batched=True (default) component draws are grouped and scattered – bit-identical to the per-draw loop (batched=False) but far faster, since each component sampler owns an independent RNG.

Parameters:

  • size (Optional[int]) – Number of iid samples to draw.

  • batched (bool) – Vectorize component draws (default); set False for the per-draw loop.

Returns:

Numpy array with shape (d,) if size is None, else with shape (size, d).

Return type:

ndarray

class GaussianMixtureAccumulator(accumulators, keys=(None, None), init='kmeans++')[source]

Bases: MixtureAccumulator

EM accumulator for Gaussian mixture responsibilities and component stats.

Vectorized E-step, fused-likelihood tracking, scalar update/initialize, combine/value/from_value, and key merge/replace behavior are inherited from MixtureAccumulator. The specialization only supplies the Gaussian mixture encoder.

Parameters:

  • accumulators (Sequence[SequenceEncodableStatisticAccumulator]) – Component accumulators, typically one multivariate Gaussian accumulator per component.

  • keys (tuple[str | None, str | None]) – Optional shared-statistic keys for weights and components.

  • init (str) – Initialization strategy, one of MIXTURE_INIT_STRATEGIES.

acc_to_encoder()[source]

Return an encoder compatible with the component accumulators.

Return type:

GaussianMixtureDataEncoder

GaussianMixtureAccumulatorFactory

alias of GaussianMixtureEstimatorAccumulatorFactory

class GaussianMixtureEstimator(estimators, name=None, conj_prior_params=None, suff_stat=None, pseudo_count=None, keys=(None, None), robust=False, init=None)[source]

Bases: MixtureEstimator

EM estimator for a multivariate Gaussian mixture.

The weight M-step and robustness options are inherited from MixtureEstimator. The specialization repacks fitted Gaussian components into (mu, sig2, w) arrays for a GaussianMixtureDistribution.

Parameters:

  • estimators (Sequence[ParameterEstimator]) – Component estimators whose fitted components carry a mean vector and a covariance (mu and covar) – typically MultivariateGaussianEstimator or DiagonalGaussianEstimator. The univariate GaussianEstimator is refused at construction (its fitted component has no covariance to pack); spell a univariate Gaussian mixture MixtureEstimator([GaussianEstimator(), ...]).

  • name (str | None) – Optional diagnostic name.

  • conj_prior_params (Any | None) – Reserved compatibility slot.

  • suff_stat (ndarray | None) – Optional prior component-count vector used with pseudo_count.

  • pseudo_count (float | None) – Optional smoothing mass for mixture weights.

  • keys (tuple[str | None, str | None]) – Optional shared-statistic keys for weights and components.

  • robust (bool) – Enable the robust weight floor, as on MixtureEstimator.

  • init (str | None) – Initialization strategy, one of MIXTURE_INIT_STRATEGIES; None selects the default "kmeans++".

accumulator_factory()[source]

Return a Gaussian mixture accumulator factory.

Return type:

GaussianMixtureEstimatorAccumulatorFactory

accumulatorFactory()[source]

Deprecated alias for accumulator_factory().

New code should call accumulator_factory. The camelCase name remains available for older callers and returns the same factory, but now emits a DeprecationWarning (see docs/support-policy.rst).

Return type:

GaussianMixtureEstimatorAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate a Gaussian mixture from aggregated EM sufficient statistics.

The component M-step and weight estimate reuse the generic mixture estimator path. The fitted component means and covariances are then repacked into a GaussianMixtureDistribution.

Parameters:

  • nobs (float | None) – Unused compatibility argument from ParameterEstimator.

  • suff_stat (tuple[ndarray, tuple[Any, ...]]) – (component_counts, component_suff_stats) tuple.

Returns:

Fitted Gaussian mixture distribution.

Return type:

GaussianMixtureDistribution

class GaussianMixtureDataEncoder(encoder=None)[source]

Bases: MixtureDataEncoder

Encoder for iid vector observations in a Gaussian mixture.

Parameters:

encoder (DataSequenceEncoder | None) – Optional component encoder. Defaults to MultivariateGaussianDataEncoder.

class HeterogeneousMixtureDistribution(components, w=MISSING, name=None, weights=MISSING)[source]

Bases: SequenceEncodableProbabilityDistribution

Mixture distribution with component-specific observation encoders.

Parameters:
compute_capabilities()[source]

Return compute-backend metadata shared by heterogeneous mixture components.

compute_declaration()[source]

Return the symbolic declaration for heterogeneous mixture weights and components.

density(x)[source]

Evaluate the density of the heterogeneous mixture at one observation.

See log_density() for details.

Parameters:

x (T) – (T): Single observation from heterogeneous mixture distribution. T is data type of components.

Returns:

Density at x.

Return type:

float

density_semantics()[source]

Return joined density semantics over heterogeneous mixture components.

log_density(x)[source]

Evaluate log-density of heterogeneous mixture distribution at observation x.

A K-component heterogeneous mixture has log-density,

log(p_mat(x)) = log(sum_{z=k}^{K} p_mat(x|z=k)*p_mat(z=k)),

where p_mat(x|z=k) is component-k log-density at x, and p_mat(z=k) = w[k]. A log-sum-exp is used to evaluate the sum inside the log of the right-hand side above. (See mixle.utils.vector.log_sum() for details).

Recall: p_mat(x|z=k) need only be compatible with same data type T. They are need not be the same distribution.

Parameters:

x (T) – (T): Single observation from heterogeneous mixture distribution. T is data type of components.

Returns:

Log-density at x.

Return type:

float

component_log_density(x)[source]

Evaluate component-wise log-density of heterogeneous mixture distribution at observation x.

A K-component heterogeneous mixture has log-density,

log(p_mat(x)) = log(sum_{z=k}^{K} p_mat(x|z=k)*p_mat(z=k)),

where p_mat(x|z=k) is component-k log-density at x, and p_mat(z=k) = w[k].

This function returns an ndarray[float] of length K, containing log(p_mat(x|z=k)) as its k^{th} entry.

Parameters:

x (T) – (T): Single observation from mixture distribution. T is data type of components.

Returns:

Numpy array of floats containing component-wise log-density at x.

Return type:

ndarray

posterior(x)[source]

Obtain the posterior distribution for each heterogeneous mixture component at observation x.

The posterior distribution of component ‘k’ at observation x is given by,

  1. p_mat(Z=k|x) = p_mat(x|Z=k)*p_mat(z=k) / p_mat(x),

where

  1. p_mat(x) = sum_{k=1}^{K} p_mat(x|Z=k)*p_mat(z=k) = sum_{k=1}^{K} p_mat(x|Z=k)*w[k].

This function returns an ndarray[float] of length K, containing p_mat(Z=k|x) as its k^{th} entry.

Parameters:

x (T) – (T): Single observation from heterogeneous mixture distribution. T is data type of components.

Returns:

Numpy array of floats containing posterior distribution at observation x.

Return type:

ndarray

seq_log_density(x)[source]

Vectorized evaluation of component-wise log-density for encoded sequence x.

Evaluates the log-density of each observation in the encoded sequence x (see log_density() for details).

Arg x must be a Tuple of length two containing and encoded from

HeterogeneousMixtureDataEncoder.seq_encode(data) with data type Sequence[T] for data.

x[0] (List[np.ndarray[int]]): The component ids for each distinct SequenceEncodableProbabilityDistribution

subclass.

x[1] (List[T1,T2,..Tk]): A list of sequence encodings of iid an iid observation sequence for each

‘k’ distinct SequenceEncodableProbabilityDistribution subclasses. The data type for each encoding is assumed to be of type Ti.

The returned value is an ndarray[float] with shape (sz,K), where K is the number of mixture components, and sz is the number of iid observations in the encoded sequence x.

Note: A row-wise log-sum-exp is performed for numerical stability. If a row contains a log-density value of,

-np.inf is returned for the corresponding observation value in the encoded sequence x.

Parameters:

x (tuple[list[ndarray], list[Any]]) – See above for details.

Returns:

Numpy array of floats containing the log_density of each observation in encoded sequence.

Return type:

ndarray

seq_component_log_density(x)[source]

Vectorized evaluation of component-wise log-density for encoded sequence x.

Arg x must be a Tuple of length two containing and encoded from

HeterogeneousMixtureDataEncoder.seq_encode(data) with data type Sequence[T] for data.

x[0] (List[np.ndarray[int]]): The component ids for each distinct SequenceEncodableProbabilityDistribution

subclass.

x[1] (List[T1,T2,..Tk]): A list of sequence encodings of iid an iid observation sequence for each

‘k’ distinct SequenceEncodableProbabilityDistribution subclasses. The data type for each encoding is assumed to be of type Ti.

Creates a 2-d numpy array of floats with vectorized evaluations of component_log_density() stored in the rows corresponding to an observation in encoded sequence x.

The returned value is an ndarray[float] with shape (sz,K), where K is the number of mixture components, and sz is the number of iid observations in the encoded sequence x.

Parameters:

x (tuple[list[ndarray], list[Any]]) – See above for details.

Returns:

2-d numpy array of floats having shape (sz,K), where sz is the number of iid obs in encoded sequence x, and K is the number of mixture components.

Return type:

ndarray

backend_seq_component_log_density(x, engine)[source]

Engine-neutral component log densities for heterogeneous encoded data.

Parameters:
Return type:

Any

backend_seq_log_density(x, engine)[source]

Engine-neutral heterogeneous-mixture log-density for encoded data.

Parameters:
Return type:

Any

seq_posterior(x)[source]

Vectorized evaluation of posterior of HeterogeneousMixtureDistribution for encoded sequence x.

Arg x must be a Tuple of length two containing and encoded from

HeterogeneousMixtureDataEncoder.seq_encode(data) with data type Sequence[T] for data.

x[0] (List[np.ndarray[int]]): The component ids for each distinct SequenceEncodableProbabilityDistribution

subclass.

x[1] (List[T1,T2,..Tk]): A list of sequence encodings of iid an iid observation sequence for each

‘k’ distinct SequenceEncodableProbabilityDistribution subclasses. The data type for each encoding is assumed to be of type Ti.

Vectorized evaluation the posterior of each observation in the encoded sequence x (see posterior() for details).

The returned value is an ndarray[float] with shape (sz,K), where K is the number of mixture components, and sz is the number of iid observations in the encoded sequence x. Each row contains the posterior of the corresponding encoded observation.

Note: A row-wise log-sum-exp is performed for numerical stability. If a row contains a log-density value of,

-np.inf is returned for the corresponding observation value in the encoded sequence x.

Parameters:

x (tuple[list[ndarray], list[Any]]) – See above for details.

Returns:

Numpy array of floats containing the posterior of each observation in encoded sequence.

Return type:

ndarray

sampler(seed=None)[source]

Return a sampler for this heterogeneous mixture.

Parameters:

seed (Optional[int]) – Seed to set for sampling with RandomState.

Returns:

HeterogeneousMixtureSampler object.

Return type:

HeterogeneousMixtureSampler

estimator(pseudo_count=None)[source]

Return an estimator with one child estimator per component.

Parameters:

pseudo_count (Optional[float]) – Used to inflate sufficient statistics in estimation.

Returns:

HeterogeneousMixtureEstimator object.

Return type:

HeterogeneousMixtureEstimator

decomposition()[source]

Heterogeneous mixture components split along the component axis (logsumexp responsibilities inside a shard; per-component stats SUM-reduce). Components are NOT homogeneous, so there is no stacked-parameter tensor to DTensor-shard (engine_axis=None -> host-shard executor mode).

dist_to_encoder()[source]

Return a data encoder assembled from the component distribution encoders.

Return type:

HeterogeneousMixtureDataEncoder

enumerator()[source]

Returns a HeterogeneousMixtureEnumerator iterating the union of component supports in descending mixture probability order.

Return type:

HeterogeneousMixtureEnumerator

class HeterogeneousMixtureEnumerator(dist)[source]

Bases: DistributionEnumerator

Enumerator over the deduplicated union of heterogeneous component supports.

Parameters:

dist (HeterogeneousMixtureDistribution)

class HeterogeneousMixtureSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler that draws a component id and delegates to that component sampler.

Parameters:

  • dist (HeterogeneousMixtureDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw iid samples from a heterogeneous mixture distribution.

The data type drawn from ‘comp_samplers’ is type T, corresponding to the data type of the mixture components.

If size is None, a single sample (of data type T) is drawn and returned. If size is not None, ‘size’-iid heterogeneous mixture samples are drawn and returned as a List with data type List[T]. With batched=True (default) component draws are grouped and scattered – bit-identical to the per-draw loop (batched=False) but far faster, since each component sampler owns an independent RNG.

Parameters:

  • size (Optional[int]) – Number of iid samples to draw.

  • batched (bool) – Vectorize component draws (default); set False for the per-draw loop.

Returns:

Data type T or List[T].

Return type:

Any | list[Any]

class HeterogeneousMixtureAccumulator(accumulators, keys=(None, None))[source]

Bases: SequenceEncodableStatisticAccumulator

EM accumulator for heterogeneous mixture weights and per-component statistics.

Parameters:

  • accumulators (list[SequenceEncodableStatisticAccumulator])

  • keys (tuple[str | None, str | None])

update(x, weight, estimate)[source]

Update sufficient statistics of HeterogeneousMixtureAccumulator with weighted observation.

Requires previous estimate of HeterogeneousMixtureDistribution.

Weights posterior of ‘estimate’ at x. Adds sum to comp_counts, then passes posterior[i] as weight for x into update() call of accumulator[i].

Parameters:

  • x (T) – Observation of heterogeneous mixture distribution.

  • weight (float) – Weight for observation.

  • estimate (HeterogeneousMixtureDistribution) – Previous iteration of EM estimate for HeterogeneousMixtureDistribution.

Returns:

None.

Return type:

None

initialize(x, weight, rng)[source]

Initialize heterogeneous-mixture sufficient statistics from one weighted observation.

If _init_rng is False, _acc_rng is set with rng. This is done for consistency in initialize and seq_initialize functions.

Initialize heterogeneous mixture weights with a sample from Dirichlet distribution. Each SequenceEncodableStatisticAccumulator is for the mixture components is initialized with a call to accumulator[i].initialize.

Parameters:

  • x (T) – Observation of heterogeneous mixture distribution.

  • weight (float) – Weight for observation.

  • rng (RandomState) – Used to set _acc_rng if not previously set.

Returns:

None.

Return type:

None

seq_initialize(x, weights, rng)[source]

Initialize heterogeneous-mixture sufficient statistics from encoded observations.

If _init_rng is False, _acc_rng is set with rng. This is done for consistency in initialize and seq_initialize functions.

Arg x must be a Tuple of length two containing and encoded from

HeterogeneousMixtureDataEncoder.seq_encode(data) with data type Sequence[T] for data.

x[0] (List[np.ndarray[int]]): The component ids for each distinct SequenceEncodableProbabilityDistribution

subclass.

x[1] (List[T1,T2,..Tk]): A list of sequence encodings of iid an iid observation sequence for each

‘k’ distinct SequenceEncodableProbabilityDistribution subclasses. The data type for each encoding is assumed to be of type Ti.

Vectorized implementation of initialize(), for sequence encoded x.

Parameters:

  • x (tuple[list[ndarray], list[Any]]) – See above for details.

  • weights (ndarray[float]) – Numpy array of positive valued floats.

  • rng (RandomState) – Used to set _acc_rng if not previously set.

Returns:

None.

Return type:

None

seq_update(x, weights, estimate)[source]

Vectorized update of sufficient statistics from encoded sequence of observations x.

Arg x must be a Tuple of length two containing and encoded from

HeterogeneousMixtureDataEncoder.seq_encode(data) with data type Sequence[T] for data.

x[0] (List[np.ndarray[int]]): The component ids for each distinct SequenceEncodableProbabilityDistribution

subclass.

x[1] (List[T1,T2,..Tk]): A list of sequence encodings of iid an iid observation sequence for each

‘k’ distinct SequenceEncodableProbabilityDistribution subclasses. The data type for each encoding is assumed to be of type Ti.

Note: Requires a previous estimate of HeterogeneousMixtureDistribution be passed. This may require seq_initialize() to be invoked prior to performing seq_update() calls.

Seq_update is similar to HeterogeneousMixtureDistribution.seq_posterior(). Results are aggregated to comp_counts and accumulators.

Parameters:

  • x (tuple[list[ndarray], list[Any]]) – See above for details.

  • weights (np.ndarray[float]) – Numpy array of positive floats.

  • estimate (MixtureDistribution) – Previous heterogeneous-mixture estimate from EM.

Returns:

None.

Return type:

None

seq_update_engine(x, weights, estimate, engine)[source]

Engine-resident mixture E-step (numpy or torch).

Component log-densities are scored on the active engine and combined into the posterior responsibilities with an engine softmax; per-component responsibilities are fed to the child accumulators. Mirrors seq_update (heterogeneous components keep per-type encodings).

combine(suff_stat)[source]

Merge the sufficient statistics of suff_stat with HeterogeneousMixtureAccumulator instance.

Arg suff_stat is a Tuple of length two containing,

suff_stat[0] (ndarray[float]): Aggregated component counts, suff_stat[1] (Tuple[T1,…,Tk]): Tuple of K sufficient statistics for the heterogeneous mixture components.

Note: The components k^{th} heterogeneous mixture component is assumed to have sufficient statistics of type Tk.

Parameters:

suff_stat (tuple[ndarray, tuple[Any, ...]]) – See above for details.

Returns:

HeterogeneousMixtureAccumulator object.

Return type:

HeterogeneousMixtureAccumulator

value()[source]

Returns sufficient statistics of MixtureAccumulator instance.

The sufficient statistics value returned (suff_stat) is a Tuple of length two containing,

suff_stat[0] (ndarray[float]): Aggregated component counts, suff_stat[1] (Tuple[T1,…,Tk]): Tuple of K sufficient statistics for the heterogeneous mixture components.

Note: The components k^{th} heterogeneous mixture component is assumed to have sufficient statistics of type Tk.

Returns:

Tuple[np.ndarray[float], Tuple[T1,…,Tk]] as described above.

Return type:

tuple[ndarray, tuple[Any, …]]

from_value(x)[source]

Set sufficient statistics of HeterogeneousMixtureAccumulator instance to x.

The sufficient statistics value ‘x’ is a Tuple of length two containing,

x[0] (ndarray[float]): Aggregated component counts, x[1] (Tuple[T1,…,Tk]): Tuple of K sufficient statistics for the heterogeneous mixture components.

Note: The components k^{th} heterogeneous mixture component is assumed to have sufficient statistics of type Tk.

Parameters:

x (tuple[ndarray, tuple[Any, ...]]) – See above for details.

Returns:

HeterogeneousMixtureAccumulator object.

Return type:

HeterogeneousMixtureAccumulator

key_merge(stats_dict)[source]

Merge this accumulator into keyed sufficient statistics.

Parameters:

stats_dict (Dict[str, Any]) – Dict mapping keys to sufficient statistics.

Returns:

None.

Return type:

None

key_replace(stats_dict)[source]

Replace this accumulator’s statistics from matching keyed values.

Parameters:

stats_dict (Dict[str, Any]) – Dict mapping keys to sufficient statistics.

Returns:

None.

Return type:

None

acc_to_encoder()[source]

Return a data encoder assembled from the component accumulator encoders.

Return type:

HeterogeneousMixtureDataEncoder

class HeterogeneousMixtureAccumulatorFactory(factories, dim, keys=(None, None))[source]

Bases: StatisticAccumulatorFactory

Factory for heterogeneous-mixture accumulators.

Parameters:

  • factories (list[StatisticAccumulatorFactory])

  • dim (int)

  • keys (tuple[str | None, str | None])

make()[source]

Return a new heterogeneous-mixture accumulator.

Return type:

HeterogeneousMixtureAccumulator

class HeterogeneousMixtureEstimator(estimators, fixed_weights=None, suff_stat=None, pseudo_count=None, name=None, keys=(None, None))[source]

Bases: ParameterEstimator

Estimator for heterogeneous mixture weights and component distributions.

Parameters:
accumulator_factory()[source]

Return an accumulator factory built from the component estimator factories.

Return type:

HeterogeneousMixtureAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate a heterogeneous mixture from aggregated sufficient statistics.

suff_stat is a two-item tuple containing:

suff_stat[0] (np.ndarray): Sufficient statistics for mixture weights. suff_stat[1] (Tuple[T1,…,Tk]): Tuple of K sufficient statistics for the heterogeneous mixture components.

suff_stat[1] is passed to the corresponding component estimator.

If fixed_weights is not None, suff_stat[0] is ignored and the estimated distribution uses the fixed weights.

If pseudo_count is set, observed counts are blended with either the estimator’s prior statistics or a uniform pseudo-count allocation.

If pseudo_count is None, observed component counts determine the weights.

Parameters:

  • nobs (Optional[float]) – Not used. Kept for consistency with ParameterEstimator super class.

  • suff_stat (tuple[ndarray, tuple[Any, ...]]) – See above for details.

Returns:

Estimated mixture distribution.

Return type:

HeterogeneousMixtureDistribution

class HeterogeneousMixtureDataEncoder(encoders)[source]

Bases: DataSequenceEncoder

Encoder that groups observations by compatible component encoder.

Parameters:

encoders (list[DataSequenceEncoder])

row_count(x)[source]

Return the shared row count from one grouped component encoding.

Parameters:

x (tuple[list[ndarray], list[Any]])

Return type:

int

seq_encode(x)[source]

Encode a sequence of iid heterogeneous mixture observations.

Note: The data type for every encoder in the keys of HeterogeneousMixtureDataEncoder attribute self.encoder_dict.keys() is T.

The returned tuple contains:

tag_list (List[ndarray[int]): Heterogeneous mixture component ids for encoded sequences in enc_data list.
enc_data (List[S1,...,Sm]): A list of 'm' encoded sequences of type Sm, corresponding to component ids
    in tag_list.
Parameters:

x (Sequence[T]) – A Sequence of iid observations drawn from a heterogeneous mixture distribution.

Returns:

Tuple[List[ndarray[int], List[S1,…,Sm]] as defined above.

Return type:

tuple[list[ndarray], list[Any]]

class PCFGTerminationCertificate(proper, reachable_nonterminals, progeny_spectral_radius, critical_singular_classes, reason)[source]

Bases: object

Finite certificate that the start-symbol branching process dies out almost surely.

Parameters:
class HeterogeneousPCFGStatistics(terminal_counts, binary_counts, emission_statistics, observation_mass, schema_version=1)[source]

Bases: Sequence[Any]

Mass-aware PCFG statistics with backward three-slot iteration.

Parameters:
class HeterogeneousPCFGDistribution(binary_rules, terminal_rules, start=None, nonterminals=None, name=None)[source]

Bases: SequenceEncodableProbabilityDistribution

CNF PCFG whose terminal productions emit arbitrary distributions.

Parameters:

  • binary_rules (BinaryRuleInput | None) – Either {parent: [(left, right, prob), ...]} or a flat sequence of (parent, left, right, prob) rules.

  • terminal_rules (TerminalRuleInput) – Either {parent: [(emission_dist, prob), ...]} or a flat sequence of (parent, emission_dist, prob) rules.

  • start (Any | None) – Start nonterminal. If omitted, the first nonterminal is used.

  • nonterminals (Sequence[Any] | None) – Optional explicit nonterminal order.

  • name (str | None) – Optional model name.

Rule probabilities are normalized over all rules sharing a parent.

density(x)[source]

Return the probability density or mass at a single observation.

Parameters:

x (Sequence[Any])

Return type:

float

log_density(x)[source]

Return the log-density or log-mass at a single observation.

Parameters:

x (Sequence[Any])

Return type:

float

seq_log_density(x)[source]

Return vectorized log-density values for sequence-encoded observations.

Parameters:

x (tuple[ndarray, tuple[Any, ...]])

Return type:

ndarray

compute_capabilities()[source]

Engine readiness intersected from the terminal emission distributions.

The CKY inside dynamic program is expressed in ComputeEngine ops (see backend_seq_log_density), so the grammar is engine-ready on whatever engines all of its terminal emission leaves support (numpy plus, e.g., torch for autograd/GPU).

backend_seq_log_density(x, engine)[source]

Engine-routed CKY inside scoring (numpy + torch).

Terminal log-densities come from each emission leaf’s own engine backend score, and the inside dynamic program runs in ComputeEngine ops. Per-sequence charts are still a Python loop (variable-length parses), so this trades the tuned NumPy _inside for engine portability and differentiability rather than raw speed.

Parameters:

x (tuple[ndarray, tuple[Any, ...]])

Return type:

Any

to_fisher(**kwargs)[source]

Inside-outside Fisher view for the PCFG.

sampler(seed=None)[source]

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

HeterogeneousPCFGSampler

enumerator()[source]

Return an exact support enumerator for acyclic enumerable PCFGs.

Return type:

HeterogeneousPCFGEnumerator

quantized_index(max_bits, bin_width_bits=1.0)[source]

Build a bounded index with an inside-style DP over quantized derivation costs.

Terminal emission distributions must themselves support quantized_index. The DP counts grammar derivations by additive quantized bit cost and lazily unranks those derivations into emitted terminal sequences. For unambiguous grammars this is an index over distinct sequences; for ambiguous grammars the same sequence can be represented by multiple derivations, and the returned log probability is still the exact PCFG log_density of that sequence.

Parameters:
Return type:

LazyQuantizedEnumerationIndex

estimator(pseudo_count=None)[source]

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

HeterogeneousPCFGEstimator

dist_to_encoder()[source]

Return the data encoder used by this distribution for vectorized methods.

Return type:

HeterogeneousPCFGDataEncoder

class HeterogeneousPCFGEnumerator(dist)[source]

Bases: DistributionEnumerator

Exact best-first enumerator for acyclic heterogeneous PCFGs.

Parameters:

dist (HeterogeneousPCFGDistribution)

class HeterogeneousPCFGSampler(dist, seed=None, max_depth=100, max_steps=10000)[source]

Bases: DistributionSampler

Sampler for HeterogeneousPCFGDistribution.

Parameters:

  • dist (HeterogeneousPCFGDistribution)

  • seed (int | None)

  • max_depth (int)

  • max_steps (int)

sample(size=None, *, batched=True)[source]

Draw one sequence, or size independent sequences, from the grammar.

Parameters:
Return type:

list[Any] | list[list[Any]]

class HeterogeneousPCFGAccumulator(emission_accumulators, num_binary_rules, terminal_parents, binary_parents, keys=(None, None))[source]

Bases: SequenceEncodableStatisticAccumulator

Inside-outside sufficient-statistic accumulator.

Parameters:

  • emission_accumulators (Sequence[SequenceEncodableStatisticAccumulator])

  • num_binary_rules (int)

  • terminal_parents (Sequence[int])

  • binary_parents (Sequence[int])

  • keys (tuple[str | None, str | None] | None)

update(x, weight, estimate)[source]

Update sufficient statistics for one observed sequence.

Parameters:
Return type:

None

initialize(x, weight, rng)[source]

Randomly initialize sufficient statistics for one observed sequence.

Parameters:
Return type:

None

seq_initialize(x, weights, rng)[source]

Randomly initialize rule and emission statistics for encoded sequences.

Parameters:
Return type:

None

seq_update(x, weights, estimate)[source]

Update encoded-sequence statistics with inside-outside posteriors.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge another accumulator value into this accumulator.

Parameters:

suff_stat (tuple[ndarray, ndarray, Sequence[Any]])

Return type:

HeterogeneousPCFGAccumulator

value()[source]

Return rule counts and emission sufficient statistics.

Return type:

HeterogeneousPCFGStatistics

from_value(x)[source]

Replace this accumulator from a serialized sufficient-statistic value.

Parameters:

x (tuple[ndarray, ndarray, Sequence[Any]])

Return type:

HeterogeneousPCFGAccumulator

scale(c)[source]

Scale rule, emission, and outer-observation statistics.

Parameters:

c (float)

Return type:

HeterogeneousPCFGAccumulator

key_merge(stats_dict)[source]

Merge keyed rule and emission statistics into stats_dict.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)[source]

Replace keyed rule and emission statistics from stats_dict.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()[source]

Return an encoder compatible with this accumulator’s emissions.

Return type:

HeterogeneousPCFGDataEncoder

class HeterogeneousPCFGAccumulatorFactory(factories, num_binary_rules, terminal_parents, binary_parents, keys=(None, None))[source]

Bases: StatisticAccumulatorFactory

Factory for HeterogeneousPCFGAccumulator.

Parameters:

  • factories (Sequence[StatisticAccumulatorFactory])

  • num_binary_rules (int)

  • terminal_parents (Sequence[int])

  • binary_parents (Sequence[int])

  • keys (tuple[str | None, str | None] | None)

make()[source]

Create a fresh heterogeneous PCFG accumulator.

Return type:

HeterogeneousPCFGAccumulator

class HeterogeneousPCFGEstimator(binary_rules, terminal_rules, start=None, nonterminals=None, pseudo_count=None, name=None, keys=(None, None))[source]

Bases: ParameterEstimator

Estimator for a fixed-topology heterogeneous PCFG.

Parameters:

  • binary_rules (BinaryRuleInput | None)

  • terminal_rules (TerminalRuleInput)

  • start (Any | None)

  • nonterminals (Sequence[Any] | None)

  • pseudo_count (float | None)

  • name (str | None)

  • keys (tuple[str | None, str | None] | None)

accumulator_factory()[source]

Return an accumulator factory for this fixed grammar topology.

Return type:

HeterogeneousPCFGAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate rule probabilities and terminal emissions from statistics.

Parameters:
Return type:

HeterogeneousPCFGDistribution

class InducedHeterogeneousPCFGEstimator(max_nonterminals, terminal_estimators, start='S', nonterminal_prefix='NT', terminal_rule_mass=0.5, rule_pseudo_count=1.0e-3, prune_threshold=0.0, min_rule_prob=0.0, name=None, keys=(None, None))[source]

Bases: ParameterEstimator

Overcomplete sparse PCFG structure learner.

This estimator builds a finite grammar skeleton automatically:

  • K nonterminals named start, prefix1, ..., prefixK-1.

  • every binary rule A -> B C.

  • every terminal rule A -> terminal_family_j.

EM learns rule probabilities and terminal emission parameters. Rules whose expected count is below prune_threshold or whose normalized probability is below min_rule_prob are assigned probability zero, but the rule layout is kept stable so repeated calls to seq_estimate remain compatible.

Parameters:

  • max_nonterminals (int)

  • terminal_estimators (Sequence[ParameterEstimator])

  • start (Any)

  • nonterminal_prefix (str)

  • terminal_rule_mass (float)

  • rule_pseudo_count (float | None)

  • prune_threshold (float)

  • min_rule_prob (float)

  • name (str | None)

  • keys (tuple[str | None, str | None] | None)

accumulator_factory()[source]

Return an accumulator factory for the induced grammar skeleton.

Return type:

HeterogeneousPCFGAccumulatorFactory

initial_model(terminal_distributions, rng=None, jitter=0.0)[source]

Create an overcomplete starting grammar without hand-written rules.

terminal_distributions may contain one distribution per terminal family, in which case each nonterminal gets its own copy of the same family list, or one distribution per generated terminal rule.

Parameters:
Return type:

HeterogeneousPCFGDistribution

estimate(nobs, suff_stat)[source]

Estimate a sparse grammar from rule counts and emission statistics.

Parameters:
Return type:

HeterogeneousPCFGDistribution

class HeterogeneousPCFGDataEncoder(terminal_encoders)[source]

Bases: DataSequenceEncoder

Encode iid PCFG observations by flattening terminal observations.

Parameters:

terminal_encoders (Sequence[DataSequenceEncoder])

seq_encode(x)[source]

Flatten token sequences and encode them for every terminal family.

Parameters:

x (Sequence[Sequence[Any]])

Return type:

tuple[ndarray, tuple[Any, …]]

row_count(x)[source]

Return the number of encoded grammar observations, not flattened terminals.

Parameters:

x (tuple[ndarray, tuple[Any, ...]])

Return type:

int

class HeterogeneousPCFGFisherView(dist)[source]

Bases: FixedFisherView

Inside-outside Fisher view for heterogeneous PCFGs.

Coordinates are expected complete-data rule counts followed by terminal emission sufficient statistics gated by the posterior probability of the corresponding terminal rule at each token. For finite enumerable PCFGs, the model Fisher is the exact observed Fisher covariance of these posterior-expected complete-data statistics under the model distribution. Recursive or infinite-support grammars should use observed_fisher_* on data.

Parameters:

dist (Any)

fisher_information(stats=None, diagonal=False, ridge=1.0e-8, **kwargs)[source]

Return model Fisher information, falling back to observed statistics.

Parameters:
Return type:

ndarray

fisher_vectors(stats=None, metric='diagonal', center=None, fisher=None, ridge=1.0e-8, **kwargs)[source]

Whiten PCFG statistics under the requested Fisher metric.

Parameters:
Return type:

ndarray

class HiddenAssociationStatistics(conditional, given, size, conditional_mass, observation_mass, schema_version=1)[source]

Bases: Sequence[Any]

Mass-aware statistics with backward three-child iteration.

Parameters:

  • conditional (Any)

  • given (Any)

  • size (Any)

  • conditional_mass (float)

  • observation_mass (float)

  • schema_version (int)

class HiddenAssociationDistribution(cond_dist, given_dist=NullDistribution(), len_dist=NullDistribution(), name=None, keys=(None, None))[source]

Bases: SequenceEncodableProbabilityDistribution

Hidden association model: values of a second set are emitted conditionally on values drawn from a first set.

Parameters:

  • cond_dist (ConditionalDistribution)

  • given_dist (SequenceEncodableProbabilityDistribution | None)

  • len_dist (SequenceEncodableProbabilityDistribution | None)

  • name (str | None)

  • keys (tuple[str | None, str | None] | None)

compute_capabilities()[source]

Return backend capability metadata for this concrete hidden association model.

compute_declaration()[source]

Return the generated-compute declaration for hidden association.

log_density(x)[source]

Log-density of the hidden association model at observation x.

For each emitted value in x[1], marginalizes the conditional emission density over the given values in x[0] weighted by their counts, then adds the log-density of the given set under given_dist and of the total emission count under len_dist.

Parameters:

x (Tuple[List[Tuple[T, float]], List[Tuple[T, float]]]) – Grouped-count observation ([(given value, count)], [(emitted value, count)]).

Returns:

Log-density at observation x.

Return type:

float

seq_log_density(x)[source]

Evaluation of log-density at sequence encoded input x (loops over log_density).

Parameters:

x (List[Tuple[List[Tuple[T, float]], List[Tuple[T, float]]]]) – Sequence encoded observations from HiddenAssociationDataEncoder.seq_encode() (the observations themselves).

Returns:

Numpy array of log-density values, one per observation.

Return type:

ndarray

backend_seq_log_density(x, engine)[source]

Evaluate encoded log-densities through distribution-owned backend composition.

Parameters:
Return type:

Any

emission_mixture(s1)[source]

The per-emission distribution q(.|S1) as a mixture, or None for an empty/degenerate S1.

q(emitted|S1) = sum_u (c_u/n1) P(emitted|u) is a finite mixture of the conditional emission distributions cond_dist.dmap[u] weighted by the given-bag’s normalized counts – enumerable whenever those component distributions are. Requires cond_dist to be a ConditionalDistribution (so the per-given components are available).

Parameters:

s1 (list[tuple[T, float]])

Return type:

MixtureDistribution | None

enumerator()[source]

Enumerate (S1, S2) observations in descending probability order.

The model factors as given_dist(S1) * [prod_e q(e|S1)^{c_e}] * P_len(n): the emitted bag S2 is drawn iid from the per-given mixture q(.|S1) (see emission_mixture()). Enumeration is a conditional product – the outer stream enumerates S1 from given_dist and, for each S1, the inner stream enumerates S2 as a multiset best-first search over q(.|S1)’s own enumeration under len_dist, merged by descending total score with given_dist(S1) as the frontier bound. Requires an enumerable non-null given_dist and a ConditionalDistribution cond_dist.

Return type:

DistributionEnumerator

sampler(seed=None)[source]

Return a sampler for iid grouped-count pairs.

Requires non-null given_dist and len_dist.

Parameters:

seed (int | None) – Optional random seed.

Returns:

A configured HiddenAssociationSampler.

Return type:

HiddenAssociationSampler

estimator(pseudo_count=None)[source]

Return an estimator built from the child estimators.

Parameters:

pseudo_count (float | None) – Unused; accepted for protocol consistency.

Returns:

A configured HiddenAssociationEstimator.

Return type:

HiddenAssociationEstimator

dist_to_encoder()[source]

Return an encoder for grouped-count hidden-association observations.

Return type:

HiddenAssociationDataEncoder

class HiddenAssociationEnumerator(dist)[source]

Bases: DistributionEnumerator

Enumerate hidden-association support pairs in best-first order.

Parameters:

dist (HiddenAssociationDistribution)

class HiddenAssociationSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for grouped-count set pairs from a hidden-association distribution.

Parameters:

  • dist (HiddenAssociationDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw iid grouped-count observations from the hidden association model.

Parameters:

  • size (Optional[int]) – Number of observations to draw. If None, a single observation is returned.

  • batched (bool)

Returns:

A ([(given value, count)], [(emitted value, count)]) tuple if size is None, else a list of such tuples of length size.

Return type:

Sequence[tuple[list[tuple[Any, float]], list[tuple[Any, float]]]] | tuple[list[tuple[Any, float]], list[tuple[Any, float]]]

sample_given(x)[source]

Draw an emitted grouped-count set conditioned on the given set x.

Parameters:

x (List[Tuple[T, float]]) – Given set as (value, count) pairs.

Returns:

List of (emitted value, count) pairs.

class HiddenAssociationAccumulator(cond_acc, given_acc=NullAccumulator(), size_acc=NullAccumulator(), name=None, keys=(None, None))[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulator for hidden-association sufficient statistics.

Parameters:

  • cond_acc (ConditionalDistributionAccumulator)

  • given_acc (SequenceEncodableStatisticAccumulator | None)

  • size_acc (SequenceEncodableStatisticAccumulator | None)

  • name (str | None)

  • keys (tuple[str | None, str | None] | None)

update(x, weight, estimate)[source]

Update sufficient statistics with the posterior assignment of emitted values to given values.

Parameters:

  • x (Tuple[List[Tuple[T, float]], List[Tuple[T, float]]]) – Grouped-count observation ([(given value, count)], [(emitted value, count)]).

  • weight (float) – Weight for the observation.

  • estimate (HiddenAssociationDistribution) – Previous estimate used to compute posteriors.

Return type:

None

initialize(x, weight, rng)[source]

Initialize sufficient statistics with random (Dirichlet) assignments of emitted to given values.

Parameters:

  • x (Tuple[List[Tuple[T, float]], List[Tuple[T, float]]]) – Grouped-count observation ([(given value, count)], [(emitted value, count)]).

  • weight (float) – Weight for the observation.

  • rng (np.random.RandomState) – Random number generator for the random assignments.

Return type:

None

seq_initialize(x, weights, rng)[source]

Initialize sufficient statistics from sequence encoded observations (loops over initialize()).

Parameters:

  • x (Sequence[Tuple[List[Tuple[T, float]], List[Tuple[T, float]]]]) – Sequence encoded observations.

  • weights (np.ndarray) – Weights, one per observation.

  • rng (np.random.RandomState) – Random number generator for the random assignments.

Return type:

None

seq_update(x, weights, estimate)[source]

Update sufficient statistics from sequence encoded observations (loops over update()).

Parameters:

  • x (Sequence[Tuple[List[Tuple[T, float]], List[Tuple[T, float]]]]) – Sequence encoded observations.

  • weights (np.ndarray) – Weights, one per observation.

  • estimate (HiddenAssociationDistribution) – Previous estimate used to compute posteriors.

Return type:

None

seq_update_engine(x, weights, estimate, engine)[source]

Engine-resident E-step for the hidden association model.

Builds the batched cross product of (given, emitted) pairs across the whole minibatch, scores them through the conditional distribution’s engine kernel, and forms the per-emitted posterior over given values with a segmented softmax (global-max shift + index_add) on the active engine. The engine-computed posterior weights are fed to the conditional accumulator; the given/size accumulators are updated per observation. Mirrors update.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge sufficient statistics of suff_stat into this accumulator.

Parameters:

suff_stat (Tuple[SS1, Optional[SS2], Optional[SS3]]) – Conditional, given, and size suff stats.

Returns:

This HiddenAssociationAccumulator.

Return type:

HiddenAssociationAccumulator

value()[source]

Return child statistics plus conditional-token and observation mass.

Return type:

HiddenAssociationStatistics

from_value(x)[source]

Set the sufficient statistics of this accumulator from x.

Parameters:

x (Tuple[SS1, Optional[SS2], Optional[SS3]]) – Conditional, given, and size suff stats.

Returns:

This HiddenAssociationAccumulator.

Return type:

HiddenAssociationAccumulator

scale(c)[source]

Scale sufficient statistics by delegating to child accumulators.

Parameters:

c (float)

Return type:

HiddenAssociationAccumulator

key_merge(stats_dict)[source]

Merge keyed statistics of the conditional, given, and size accumulators into stats_dict.

Parameters:

stats_dict (Dict[str, Any]) – Maps keys to merged sufficient statistics.

Return type:

None

key_replace(stats_dict)[source]

Replace keyed statistics of the conditional, given, and size accumulators with those in stats_dict.

Parameters:

stats_dict (Dict[str, Any]) – Maps keys to merged sufficient statistics.

Return type:

None

acc_to_encoder()[source]

Return an encoder compatible with hidden-association observations.

Return type:

HiddenAssociationDataEncoder

class HiddenAssociationAccumulatorFactory(cond_factory, given_factory=NullAccumulatorFactory(), len_factory=NullAccumulatorFactory(), name=None, keys=(None, None))[source]

Bases: StatisticAccumulatorFactory

Factory for hidden-association accumulators.

Parameters:

  • cond_factory (ConditionalDistributionAccumulatorFactory)

  • given_factory (StatisticAccumulatorFactory | None)

  • len_factory (StatisticAccumulatorFactory | None)

  • name (str | None)

  • keys (tuple[str | None, str | None] | None)

make()[source]

Return a fresh hidden-association accumulator.

Return type:

HiddenAssociationAccumulator

class HiddenAssociationEstimator(cond_estimator, given_estimator=NullEstimator(), len_estimator=NullEstimator(), pseudo_count=None, name=None, keys=(None, None))[source]

Bases: ParameterEstimator

Estimator for hidden association distributions.

Parameters:

  • cond_estimator (ConditionalDistributionEstimator)

  • given_estimator (ParameterEstimator | None)

  • len_estimator (ParameterEstimator | None)

  • pseudo_count (float | None)

  • name (str | None)

  • keys (tuple[str | None, str | None] | None)

accumulator_factory()[source]

Return an accumulator factory matching this estimator.

Return type:

HiddenAssociationAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate a HiddenAssociationDistribution from aggregated sufficient statistics.

Parameters:

  • nobs (Optional[float]) – Number of observations, passed to the given and length estimators.

  • suff_stat (Tuple[SS1, Optional[SS2], Optional[SS3]]) – Conditional, given, and size suff stats.

Returns:

HiddenAssociationDistribution object.

Return type:

HiddenAssociationDistribution

class HiddenAssociationDataEncoder[source]

Bases: DataSequenceEncoder

Encode grouped-count hidden-association observations.

seq_encode(x)[source]

Encode a sequence of iid grouped-count observations (identity encoding).

Parameters:

x (Sequence[Tuple[List[Tuple[T, float]], List[Tuple[T, float]]]]) – Sequence of iid ([(given value, count)], [(emitted value, count)]) observations.

Returns:

The observations unchanged (seq_log_density and seq_update loop over them).

Return type:

Sequence[tuple[list[tuple[T, float]], list[tuple[T, float]]]]

row_count(x)[source]

Return the number of identity-encoded grouped-count observations.

Parameters:

x (Any)

Return type:

int

hmm_dirichlet_default_prior(num_states)[source]

Returns the default (init_prior, row_priors) pair of unit-parameter Dirichlets.

Parameters:

num_states (int) – Number of hidden states S.

Returns:

Tuple (DirichletDistribution, list of S DirichletDistribution).

terminal_forward_loglik(log_w, log_a, log_b, term_mask)[source]

Log-likelihood of a terminal-state HMM sequence (length is a stopping time at the first terminal state).

log_b is the per-position, per-state emission log-density (L, K). The forward only transitions from non-terminal states and the likelihood sums the final position over terminal states. Shared by every HMM variant whose forward is a linear state trellis (base, lookback, …).

Parameters:
Return type:

float

terminal_forward_backward(log_w, log_a, log_b, term_mask)[source]

Terminal-state forward-backward; returns (loglik, gamma (L,K), xi (L-1,K,K)) (gamma/xi None if 0-prob).

The backward pass mirrors the forward: only the final position may be terminal, and only non-terminal states have a future. Responsibilities are normalized by the sequence likelihood.

Parameters:
Return type:

tuple[float, ndarray | None, ndarray | None]

class HiddenMarkovModelDistribution(topics=MISSING, w=MISSING, transitions=MISSING, taus=None, len_dist=NullDistribution(), name=None, terminal_values=None, use_numba=None, weights=MISSING, prior=None, terminal_states=None, components=MISSING)[source]

Bases: SequenceEncodableProbabilityDistribution

Hidden Markov model distribution for variable-length observation sequences.

Parameters:

  • topics (Sequence[SequenceEncodableProbabilityDistribution])

  • w (Sequence[float] | np.ndarray)

  • transitions (list[list[float]] | np.ndarray)

  • taus (list[list[float]] | np.ndarray | None)

  • len_dist (SequenceEncodableProbabilityDistribution | None)

  • name (str | None)

  • terminal_values (set[T] | None)

  • use_numba (bool | None)

  • weights (Sequence[float] | np.ndarray)

  • terminal_states (set[int] | Sequence[int] | None)

  • components (Sequence[SequenceEncodableProbabilityDistribution])

get_prior()[source]

Returns the chain conjugate prior in (init_prior, row_priors) form (or None).

Per-state emission component priors are owned by the emission distributions themselves.

set_prior(prior)[source]

Set the conjugate Dirichlet chain prior and precompute its digamma expectations.

With Dirichlet init_prior and Dirichlet row_priors (over the fixed hidden states 0..S-1) this caches the digamma expectations E[ln p_k] = psi(alpha_k) - psi(sum alpha) used by expected_log_density and sets has_conj_prior accordingly. prior=None leaves the distribution a plain point model.

Parameters:

prior – (init_prior, row_priors) tuple or None.

Return type:

None

expected_log_density(x)[source]

Forward log-likelihood with digamma-expected initial/transition log-probabilities and the topics’ expected_log_density emissions.

Falls back to the plug-in log_density(x) when no conjugate prior is set. Not supported for the taus/topic-mixture parameterization (falls back to log_density there).

Parameters:

x (List[T]) – Observed sequence of HMM emissions.

Returns:

Expected log-density of the observed HMM sequence x.

Return type:

float

seq_expected_log_density(x)[source]

Vectorized expected_log_density() at sequence-encoded input x.

Falls back to seq_log_density(x) when no conjugate prior is set or for the taus parameterization.

Parameters:

x (tuple[tuple[int, list[tuple[int, int]], list[ndarray], ndarray, ndarray, ndarray, Any], Any, Any | None] | tuple[tuple[ndarray, ndarray, ndarray], Any | None]) – Encoded sequences from seq_encode().

Returns:

Numpy array of expected log-densities, one per sequence.

Return type:

ndarray

compute_capabilities()[source]

Return compute-backend metadata for HMM scoring and Baum-Welch updates.

compute_declaration()[source]

Return the symbolic declaration for HMM chain, emission, and length statistics.

density(x)[source]

Returns the density of HMM for an observed sequence x.

See ‘HiddenMarkovDistribution.log_density()’ for details.

Parameters:

x (List[T]) – Observed sequence of HMM emissions.

Returns:

Density of HMM for observed sequence x.

Return type:

float

log_density(x)[source]

Returns the log-density of HMM for observed sequence x.

Density for a sequence of length N is given by recursively evaluating the conditional density,

p_mat(x_mat(0),x_mat(1),....,x_mat(t)) = p_mat(x_mat(t) given x_mat(0),...,x_mat(t-1)) = p_mat(x_mat(t) given Z(t)) * p_mat(Z(t) given Z(t-1)) * p_mat(Z(t-1) given x_mat(0),....,x_mat(t-1))

for t = 1,2,…,N-1. p_mat(Z(0)) is given by ‘w’, p_mat(x_mat(t)|Z(t)) is given by emission distribution ‘topics’ for t = 0,1,…,N-1.

The returned density is given by

p_mat(x_mat) = p_mat(x_mat(0),x_mat(1),….,x_mat(t))*P_len(N).

where P_len(N) is the length distribution ‘len_dist’, if assigned. Note: All calculations are done on the log scale with log-sum-exp used to prevent numerical underflow.

If ‘has_topics’ is true, ‘weighed_log_sum_exp’ and ‘log_sum’ calls from mixle.utils.vector are used to handle the emission distributions being treated as mixture distributions with weights ‘log_taus’.

Parameters:

x (List[T]) – Observed sequence of HMM emissions.

Returns:

Log-density of observed HMM sequence x.

Return type:

float

seq_log_density(x)[source]

Return one HMM log-density value per encoded observation sequence.

Impossible observation/state combinations are reported as -inf in the returned vector. Valid finite inputs should not produce NaN in the public result.

Parameters:

x (tuple[tuple[int, list[tuple[int, int]], list[ndarray], ndarray, ndarray, ndarray, Any], Any, Any | None] | tuple[tuple[ndarray, ndarray, ndarray], Any | None])

Return type:

ndarray

backend_seq_log_density(x, engine)[source]

Engine-neutral forward scores for non-numba encoded HMM batches.

The compiled/numba encoding remains on the legacy NumPy path. The standard blocked encoding is converted through the active engine and composes child distribution-owned backend scores.

Parameters:
Return type:

Any

seq_posterior(x, filtered=False)[source]

Return per-sequence posterior state probabilities for encoded observations.

Each returned array holds the forward-backward SMOOTHING marginals P(z_t = k | x_1..T) (one row per position), matching every other seq_posterior in the package. Pass filtered=True for the forward-only filtered probabilities P(z_t = k | x_1..t).

A base install has no numba (it is the [numba] extra), and terminal_states turns the numba route off outright, so this used to return None on the default configuration – a documented API that silently produced nothing (P02-F01). Those two routes now decode the batch into per-sequence emission log-densities and run the same exact forward-backward latent_posterior uses, with the terminal-state restriction folded into the last position’s emissions.

Parameters:
Return type:

list[ndarray] | None

property components: tuple[SequenceEncodableProbabilityDistribution, ...]

The emission distributions, under the constructor’s documented components alias.

HiddenMarkovModelDistribution(components=...) has been accepted since 0.8.0 but the attribute only ever existed as topics, so the alias did not round-trip the way MixtureDistribution.components does (P02-F08).

viterbi(x)[source]

Return the most likely latent-state path for a single observation sequence.

The path is recovered by backpointer backtracking (max-product / Viterbi), so it is the jointly most probable state SEQUENCE – not the per-position argmax of the score matrix.

Under terminal_states the search is restricted to admissible paths, the same way latent_posterior() restricts the sum: a terminal state may occupy only the LAST position. This ignored the restriction and returned the unrestricted maximizer, which on a four-step chain was a path visiting a terminal state twice in the middle – a path the model assigns probability zero (R05-F01).

Parameters:

x (list[T])

Return type:

ndarray

latent_posterior(x)[source]

Return the exact chain posterior q(z | x) over hidden states for one observation sequence.

The returned MarkovChainLatentPosterior can .marginals() (forward-backward smoothing probabilities), .sample(rng) a full state path by FFBS, .mode() (the Viterbi path), or .entropy() (the exact chain entropy).

Under terminal_states the restriction is folded into the emissions rather than left to the caller: a terminal state may occupy only the LAST position, so its emission is impossible everywhere else, and the ordinary recursion over those emissions is then exactly the restricted one (checked against enumeration of the admissible paths). This ignored terminal_states entirely and returned the unrestricted model’s posterior – on a four-step chain whose last observation belongs to a non-terminal state, a marginal off by 0.999 – and with it .mode(), .sample() and posterior_predictive(), all of which read this object (R05-F01).

Parameters:

x (list[T])

Return type:

MarkovChainLatentPosterior

posterior_predictive(x, seed=None)[source]

Draw a new observation sequence conditioned on x.

Sample a full hidden-state path from the posterior q(z | x) by FFBS, then emit a fresh observation from each state’s emission distribution – “given the sequence I saw, draw a new sequence from the states it most likely passed through”. Returns a list the length of x.

Parameters:
Return type:

list[Any]

seq_viterbi(x)[source]

Return maximum-likelihood hidden-state assignments for encoded sequences.

Each sequence’s assignment is its Viterbi path, recovered by backpointer backtracking. The result is indexed in the encoder’s flattened sequence layout (time-banded for the blocked encoding, sequence-contiguous for the numba encoding). Use the data encoder metadata to map flattened assignments back to individual sequences when working directly with encoded data.

Under terminal_states each sequence’s search is restricted to admissible paths exactly as viterbi() restricts it, so the two agree sequence by sequence.

Parameters:

x (tuple[tuple[int, list[tuple[int, int]], list[ndarray], ndarray, ndarray, ndarray, Any], Any, Any | None] | tuple[tuple[ndarray, ndarray, ndarray], Any | None])

to_fisher(**kwargs)[source]

Return a Fisher view backed by HMM forward-backward statistics.

density_semantics()[source]

Return joined density semantics over emissions and the optional length distribution.

sampler(seed=None)[source]

Create a sampler for hidden paths and emitted observations.

Sampling requires either a length distribution or terminal-state stopping behavior. When a length distribution is present, it should be a sequence-encodable distribution over nonnegative integers.

Parameters:

seed (int | None) – Optional seed for the sampler’s random state.

Returns:

A HiddenMarkovSampler configured with the model parameters.

Return type:

HiddenMarkovSampler

estimator(pseudo_count=None)[source]

Create an estimator initialized from this hidden Markov distribution.

Parameters:

pseudo_count (Optional[float]) – Prior mass used to smooth transition and emission statistics.

Returns:

Estimator configured with matching emission and length estimators.

Return type:

HiddenMarkovEstimator

decomposition()[source]

The HMM splits along its STATE axis (the per-state emission distributions).

Unlike a mixture this is NOT suff-stat-separable – the forward-backward couples all states across time – so the executor does not reduce it; instead the per-state emission scoring and accumulation (the dominant cost for rich emissions) are distributed inside the Baum-Welch E-step (host-shard mode, engine_axis=None), while the recursion stays serial. Exposing the axis lets the balance planner use the cluster for a massive HMM even on a single observation sequence.

dist_to_encoder()[source]

Return an encoder for HMM observation sequences and optional lengths.

Return type:

HiddenMarkovDataEncoder

enumerator()[source]

Return an enumerator over observation sequences in descending marginal probability order.

Return type:

HiddenMarkovModelEnumerator

determinize(max_states=1 << 16, max_denominator=10**9)[source]

Weighted determinization (Mohri 1997; Mohri & Riley 2002) of this terminal-value HMM into a DeterminizedSequenceDistribution.

Rebuilds the (possibly ambiguous) machine over belief states so each sequence has a single path and edge weights multiply to the exact marginal – giving exact, duplicate-free n-best sequences and sub-linear structural seek, where ranking the original HMM gives n-best paths. Float probabilities are rationalized (max_denominator) for decidable belief-equality. Requires terminal_values and finite/enumerable emissions; raises EnumerationError if not finitely determinizable within max_states (the twins property fails – keep the original HMM’s exact O(index) path instead).

Parameters:

  • max_states (int)

  • max_denominator (int)

quantized_count_index(quantizer, max_fine_bucket)[source]

BoundedCount for the MARGINAL HMM law: a forward count DP over the trellis with an iterative emission-split unrank, reaching a 2**M budget structurally.

log p(x) = logsumexp over latent paths. We count (state-path, observation) PAIRS by their joint cost log w_{s0} + sum_t log trans(s_t|s_{t-1}) + sum_t log emit_{s_t}(x_t): the forward DP pools paths into a per-(length, end-state) count histogram, where each step convolves the prefix histogram with the emission’s count index (choosing the emitted symbol). This is the HMM analogue of the Mixture bound – a conservative UPPER bound that does NOT deduplicate an observation produced by multiple paths and bins by the joint (dominant-path / tropical) cost rather than the exact logsumexp; every unranked value still carries its exact marginal log_density. Unranking is one iterative backward walk over t: each step does a local times-split (recover the emitted symbol + bucket split) and a plus-choice (recover the predecessor state) – O(L), no recursion. Falls back to capped enumerate-and-bin for non-plain HMMs (taus / terminal_values) or emissions that cannot count structurally.

Parameters:

max_fine_bucket (int)

is_canonical_copy(value, coarse_bin, quantizer)[source]

Stateless dedup: keep an observation only at its min-cost (canonical) path’s bin.

The structural index emits an observation once per state-path that can generate it; the canonical copy is the one at the minimal joint fine bucket. A min-plus forward pass over the trellis computes that minimum exactly (mirroring the count-index’s fine-bucket sums), so the check is O(L * n_states^2) with no state. Falls back to True for non-plain HMMs.

Parameters:

coarse_bin (int)

Return type:

bool

class HiddenMarkovModelEnumerator(dist, topics=None, log_w=None, log_transitions=None, len_dist=None, path_root=None)[source]

Bases: DistributionEnumerator

Best-first enumerator over HMM observation sequences marginalized over hidden states.

Parameters:

  • dist (SequenceEncodableProbabilityDistribution)

  • topics (Sequence[SequenceEncodableProbabilityDistribution] | None)

  • log_w (np.ndarray | None)

  • log_transitions (np.ndarray | None)

  • len_dist (SequenceEncodableProbabilityDistribution | None)

  • path_root (str | None)

class HiddenMarkovSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for finite-length or terminal-value HMM observation sequences.

Parameters:

  • dist (HiddenMarkovModelDistribution)

  • seed (int | None)

sample_seq(size=None, *, batched=True)[source]

Sample iid HMM sequences from the length-distribution path.

If size is None, one sequence is returned. Otherwise the result is a list of size sampled sequences.

With batched=True (default) the hidden-state paths for the whole batch are drawn in a single vectorized pass (the MarkovChainSampler advances all chains across time at once) and the emissions are drawn by grouping token positions by hidden state and invoking each emission sampler once. Emission batching is byte-identical to the legacy nested loop (each emission sampler owns an independent RandomState consumed in state order), but vectorizing the state path changes the RNG consumption order, so the state paths (and therefore the emissions conditioned on them) are only statistically equivalent to batched=False – not byte-identical. Set batched=False to reproduce the exact legacy output for a given seed.

Parameters:

  • size (int | None) – Number of iid HMM sequences to sample.

  • batched (bool) – Vectorize state-path and emission draws (default); set False for the legacy per-draw loop.

Returns:

One sequence or a list of sampled sequences, depending on size.

Return type:

list[Any] | list[list[Any]]

sample_terminal(terminal_set, max_steps=DEFAULT_TERMINAL_STEP_CAP)[source]

Sample one HMM sequence until an emitted terminal value appears.

Parameters:

  • terminal_set (set[T]) – Values that stop the sampled sequence.

  • max_steps (int)

Returns:

A sequence whose final value belongs to terminal_set.

Return type:

list[T]

sample_terminal_states(max_steps=DEFAULT_TERMINAL_STEP_CAP)[source]

Sample an HMM sequence run until the hidden chain first enters an absorbing (terminal) state.

The path z_1, z_2, ... is drawn from the chain and stops the moment z_L is terminal; the returned sequence emits one observation per state, so its last state is terminal and all earlier states are not. cap guards against a terminal state that is unreachable.

Parameters:

max_steps (int)

Return type:

list[T]

sample(size=None, *, batched=True)[source]

Draw iid HMM observation sequences.

Length-distribution HMMs delegate to sample_seq(). Terminal-state and terminal-value HMMs sample until their stopping condition is reached.

With batched=True (default) the length-distribution path uses the vectorized sample_seq() (statistically equivalent, not byte-identical – see its docstring). The terminal-value path is inherently sequential and always uses the legacy loop. batched=False reproduces the exact legacy output for a given seed.

Parameters:

  • size (int | None) – Number of iid HMM sequences to sample.

  • batched (bool) – Vectorize state-path and emission draws on the length-distribution path (default); set False for the legacy per-draw loop.

Returns:

One sequence or a list of sampled sequences, depending on size.

class HiddenMarkovAccumulator(accumulators, len_accumulator=NullAccumulator(), use_numba=None, keys=(None, None, None), name=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Baum-Welch sufficient-statistic accumulator for HMM chain and emission parameters.

Parameters:

  • accumulators (Sequence[SequenceEncodableStatisticAccumulator])

  • len_accumulator (SequenceEncodableStatisticAccumulator | None)

  • use_numba (bool | None)

  • keys (tuple[str | None, str | None, str | None])

  • name (str | None)

update(x, weight, estimate)[source]

Update sufficient statistics of HiddenMarkovAccumulator with one observation.

Note: Not efficient. Prefer seq_encode() for fully encoded sequences.

Parameters:

  • x (List[T]) – HMM observation sequence.

  • weight (float) – Weight for observation.

  • estimate (HiddenMarkovModelDistribution) – Previous estimate of HMM.

Returns:

None.

Return type:

None

initialize(x, weight, rng)[source]

Initialize hidden-Markov sufficient statistics from one weighted sequence.

Parameters:

  • x (List[T]) – HMM observation sequence.

  • weight (float) – Weight for observation.

  • rng (RandomState) – Sets RandomState member values if not already set.

Returns:

None.

Return type:

None

seq_initialize(x, weights, rng)[source]

Vectorized initialization of HiddenMarkovAccumulator.

Note: Initialization method depends on sequence encoding for Numba or baseline numpy. Both methods do not call numba for initialization.

If _init_rng is False, protected RandomState members are set from rng for the accumulators. This ensures initialize() method produces a consistent initialization for the same datasets.

The input ‘x’ is a sequence encoded HMM sequence of iid observations produced by ‘HiddenMarkovDataEncoder.seq_encode()’. Arg x is either Tuple[None, enc] or Tuple[None, enc_numba].

For the first case, enc is Tuple[Tuple[….], T_topic, T_len], where the first tuple is given by a Tuple of

enc[0][0] (int): Total number of observed emissions from all HMM sequences.
enc[0][1] (List[Tuple[int, int]]): Contains bands for t^th observation in HMM sequences stored in 'seq_x'.
enc[0][2] (List[ndarray[int]]): List of numpy array on sequence indices that have a next observed emission.
enc[0][3] (np.ndarray[int]): Numpy array of sequence lengths.
enc[0][4] (np.ndarray[int]): 2-d matrix with rv[0][0] rows, and column length equal to the length of the
    largest HMM sequence. This is used to store the index of seq_x corresponding to emission x[i][t]. A -1
    is stored if the sequence length has already been met.
enc[0][5] (ndarray): Numpy array containing lists index 'i' corresponding to x[i][t] block of 'seq_x'.
enc[0][6] (T_topic): Sequence encoded value of 'seq_x'.

The next two entries of the Tuple are,

enc[1] (T_topic): Sequence encoded observation values in order. Just for seq_init consistency.
enc[1] (Optional[T_len]): Sequence encoded value of lengths of HMM distribution. None if len_encoder is
    the NullDataEncoder.
The first entry of enc_numba is a Tuple of length-3,

enc_numba[0][0] (ndarray[int]): Sequence id’s for observed values. enc_numba[0][1] (ndarray[int]): Sequence lengths for each observed HMM sequence. enc_numba[0][2] (T_topic): Sequence encoded observation values.

The second entry is,
enc_numba[1] (Optional[T_len]): Sequence encoded values of sequence lengths. None if len_encoder is

NullDataEncoder.

Parameters:

  • x – See above for details.

  • weights (np.ndarray) – Numpy array of weights for observations.

  • rng (RandomState) – Used to set seed on random initialization.

Returns:

None.

Return type:

None

seq_update(x, weights, estimate)[source]

Vectorized accumulator update from encoded HMM observation sequences.

This is a vectorized implementation of the Baum-Welch algorithm. If use_numba, Numba functions are called for the alpha and beta pass. Else, a vectorized Numpy implementation of Baum-Welch is used.

The input ‘x’ is a sequence encoded HMM sequence of iid observations produced by ‘HiddenMarkovDataEncoder.seq_encode()’. Arg x is either Tuple[None, enc] or Tuple[None, enc_numba].

For the first case, enc is Tuple[Tuple[….], T_topic, T_len], where the first tuple is given by a Tuple of

enc[0][0] (int): Total number of observed emissions from all HMM sequences.
enc[0][1] (List[Tuple[int, int]]): Contains bands for t^th observation in HMM sequences stored in 'seq_x'.
enc[0][2] (List[ndarray[int]]): List of numpy array on sequence indices that have a next observed emission.
enc[0][3] (np.ndarray[int]): Numpy array of sequence lengths.
enc[0][4] (np.ndarray[int]): 2-d matrix with rv[0][0] rows, and column length equal to the length of the
    largest HMM sequence. This is used to store the index of seq_x corresponding to emission x[i][t]. A -1
    is stored if the sequence length has already been met.
enc[0][5] (ndarray): Numpy array containing lists index 'i' corresponding to x[i][t] block of 'seq_x'.
enc[0][6] (T_topic): Sequence encoded value of 'seq_x'.

The next two entries of the Tuple is,

enc[1] (T_topic): Sequence encoded observation values in order. Just for seq_init consistency.
enc[2] (Optional[T_len]): Sequence encoded value of lengths of HMM distribution. None if len_encoder is
    the NullDataEncoder.
The first entry of enc_numba is a Tuple of length-3,

enc_numba[0][0] (ndarray[int]): Sequence id’s for observed values. enc_numba[0][1] (ndarray[int]): Sequence lengths for each observed HMM sequence. enc_numba[0][2] (T_topic): Sequence encoded observation values.

The second entry is,
enc_numba[1] (Optional[T_len]): Sequence encoded values of sequence lengths. None if len_encoder is

NullDataEncoder.

Parameters:

  • x – See above for details.

  • weights (np.ndarray) – Numpy array of weights for observation.

  • estimate (HiddenMarkovModelDistribution) – Previous EM estimate of HMM model.

Returns:

None.

Return type:

None

seq_update_engine(x, weights, estimate, engine)[source]

Engine-resident Baum-Welch E-step.

Mirrors seq_update() but computes the posteriors (gamma), expected transition counts (xi), and initial-state posteriors (pi) with hmm_engine_forward_backward(), so the forward-backward runs on the active engine (numpy or torch/GPU/autograd). Per-state emission statistics are accumulated by the existing child accumulators using the engine-computed gamma weights. Falls back to the host seq_update() for the blocked (non-numba) encoding.

Parameters:

estimate (HiddenMarkovModelDistribution)

Return type:

None

combine(suff_stat)[source]

Combine the sufficient statistics of HiddenMarkovAccumulator with suff_stat arg.

Sufficient statistics in suff_stat are a Tuple containing:

suff_stat[0] (int): Number of hidden states. suff_stat[1] (np.ndarray): Initial state counts. suff_stat[2] (np.ndarray): State counts. suff_stat[3] (np.ndarayy): State transition counts. suff_stat[4] (Sequence[T1]): Emission distribution accumulators. suff_stat[5] (Optional[T2]): Optional sufficient statistics of the length distribution.

Note: T1 is the assumed type for the emission accumulator sufficient statistics. T2 is the assumed type for the length accumulator sufficient statistics.

Parameters:

suff_stat (tuple[int, ndarray, ndarray, ndarray, Sequence[T1], T2 | None]) – See above for details.

Returns:

HiddenMarkovAccumulator object.

Return type:

HiddenMarkovAccumulator

value()[source]

Return accumulated hidden Markov sufficient statistics.

Returned value rv is a Tuple containing:

rv[0] (int): Number of hidden states. rv[1] (np.ndarray): Initial state counts. rv[2] (np.ndarray): State counts. rv[3] (np.ndarray): State transition counts. rv[4] (Sequence[T1]): Emission distribution accumulator sufficient statistics (type T1). rv[5] (Optional[T2]): Optional sufficient statistics of the length distribution (type T2).

Note: T1 is the assumed type for the emission accumulator sufficient statistics. T2 is the assumed type for the length accumulator sufficient statistics.

Returns:

Tuple[int, np.ndarray, np.ndarray, np.ndarray, Sequence[T1], Optional[T2]].

Return type:

tuple[int, ndarray, ndarray, ndarray, Sequence[Any], Any | None]

from_value(x)[source]

Restore accumulator state from sufficient statistics.

Returned value x is a Tuple containing:

x[0] (int): Number of hidden states. x[1] (np.ndarray): Initial state counts. x[2] (np.ndarray): State counts. x[3] (np.ndarray): State transition counts. x[4] (List[T1]): Emission distribution accumulators. x[5] (Optional[T2]): Optional sufficient statistics of the length distribution.

Note: T1 is the assumed type for the emission accumulator sufficient statistics. T2 is the assumed type for the length accumulator sufficient statistics.

Parameters:

x (tuple[int, ndarray, ndarray, ndarray, Sequence[T1], T2 | None]) – See above for details.

Returns:

This accumulator after restoration.

Return type:

HiddenMarkovAccumulator

scale(c)[source]

Scale linear HMM sufficient statistics while preserving metadata.

Parameters:

c (float)

Return type:

HiddenMarkovAccumulator

key_merge(stats_dict)[source]

Merge this accumulator into keyed sufficient statistics.

Parameters:

stats_dict (Dict[str, Any]) – Dictionary containing sufficient statistics for corresponding keys.

Returns:

None.

Return type:

None

key_replace(stats_dict)[source]

Replace keyed sufficient-statistic arrays from stats_dict.

Parameters:

stats_dict (dict[str, Any]) – Mapping from configured statistic keys to replacement initial-state, transition, emission, or length statistics.

Return type:

None

acc_to_encoder()[source]

Return an encoder matching this accumulator’s emissions, lengths, and backend.

Return type:

HiddenMarkovDataEncoder

class HiddenMarkovAccumulatorFactory(factories, len_factory=NullAccumulatorFactory(), use_numba=None, keys=(None, None, None), name=None)[source]

Bases: StatisticAccumulatorFactory

Factory for HMM Baum-Welch sufficient-statistic accumulators.

Parameters:

  • factories (Sequence[StatisticAccumulatorFactory])

  • len_factory (StatisticAccumulatorFactory)

  • use_numba (bool | None)

  • keys (tuple[str | None, str | None, str | None] | None)

  • name (str | None)

make()[source]

Return a fresh HMM accumulator with independent child accumulators.

Return type:

HiddenMarkovAccumulator

class HiddenMarkovEstimator(estimators, len_estimator=NullEstimator(), pseudo_count=(None, None), name=None, keys=(None, None, None), use_numba=None, prior=None, steady_state_init=False, terminal_states=None)[source]

Bases: ParameterEstimator

Estimator for HMM initial, transition, emission, and optional length distributions.

Parameters:

  • estimators (list[ParameterEstimator])

  • len_estimator (ParameterEstimator | None)

  • pseudo_count (float | tuple[float | None, float | None] | None)

  • name (str | None)

  • keys (tuple[str | None, str | None, str | None] | None)

  • use_numba (bool | None)

  • steady_state_init (bool)

  • terminal_states (set[int] | Sequence[int] | None)

accumulator_factory()[source]

Return a factory that accumulates Baum-Welch statistics for this estimator.

get_prior()[source]

Returns the chain conjugate prior in (init_prior, row_priors) form (or None).

Per-state emission component priors are owned by the topic estimators themselves.

set_prior(prior)[source]

Set the conjugate Dirichlet chain prior and flag whether it admits the conjugate update.

Parameters:

prior – (init_prior, row_priors) tuple or None; has_conj_prior is set when both the initial-state prior and all row priors are Dirichlet.

Return type:

None

model_log_density(model)[source]

Log-density of the model parameters under the priors (ELBO global term).

Sums the Dirichlet log-densities of the initial-state and transition probabilities (floored at a tiny constant so boundary MAP estimates score finitely) plus each topic estimator’s model_log_density of its emission distribution. Returns the emission-only sum without a conjugate chain prior.

Parameters:

model (HiddenMarkovModelDistribution) – Model to score.

Returns:

Prior log-density of the model parameters.

Return type:

float

estimate(nobs, suff_stat)[source]

Estimate HiddenMarkovModel from aggregated sufficient statistics contained in arg ‘suff_stat’.

Sufficient statistics in arg ‘suff_stat’ are a Tuple containing:

suff_stat[0] (int): Number of hidden states.
suff_stat[1] (np.ndarray): Initial state counts.
suff_stat[2] (np.ndarray): State counts.
suff_stat[3] (np.ndarayy): State transition counts.
suff_stat[4] (List[T1]): List of Sufficient statistics for the emission distribution accumulators.
    Each having type S0.
suff_stat[5] (Optional[T2]): Optional sufficient statistics of the length distribution.

Note: T1 is the type for the sufficient statistics of the emission accumulators. T2 is the type for the length accumulator.

If pseudo_count[0] is not None, the initial counts in ‘suff_stat’ is re-weighted in estimation. If pseudo_count[1] is not None, the transition counts in ‘suff_stat’ are re-weighted in estimation.

Parameters:
Returns:

HiddenMarkovModelDistribution object.

Return type:

HiddenMarkovModelDistribution

class HiddenMarkovDataEncoder(emission_encoder, len_encoder=NullDataEncoder(), use_numba=None)[source]

Bases: DataSequenceEncoder

Encoder for HMM observation sequences and optional sequence lengths.

Parameters:

  • emission_encoder (DataSequenceEncoder)

  • len_encoder (DataSequenceEncoder | None)

  • use_numba (bool | None)

seq_encode(x)[source]

Sequence encode sequences of iid HMM observations.

Numba sequence encoding: Return type Tuple[Tuple[np.ndarray, np.ndarray, T_topic], Optional[T_len]] where T_topicis the type for ‘emission_encoder.seq_encode()’ and T_len is the type for ‘len_encoder.seq_encode()’. The first entry of the returned value (rv_numba) is a Tuple of length-3,

rv_numba[0][0] (ndarray[int]): Sequence id's for observed values.
rv_numba[0][1] (ndarray[int]): Sequence lengths for each observed HMM sequence.
rv_numba[0][2] (T_topic): Sequence encoded observation values.
rv_numba[1] (Optional[T_len]): Sequence encoded values of sequence lengths. None if len_encoder is
    NullDataEncoder.

If use_numba is False, calls HiddenMarkovDataEncoder._seq_encode(x). (See ‘_seq_encode’ for details).

Parameters:

x (List[List[T]]) – A sequence of iid observations from an HMM distribution of type T.

Returns:

Tuple[None, rv_numba] if use_numba, else Tuple[rv, None].

Return type:

tuple[tuple[tuple[int, list[tuple[int, int]], list[ndarray], ndarray, ndarray, ndarray, Any], Any, Any | None] | None, tuple[tuple[ndarray, ndarray, ndarray], Any | None] | None]

row_count(x)[source]

Return and validate the number of encoded observation sequences.

Parameters:

x (Any)

Return type:

int

HiddenMarkovModelSampler

alias of HiddenMarkovSampler

HiddenMarkovModelAccumulator

alias of HiddenMarkovAccumulator

HiddenMarkovModelAccumulatorFactory

alias of HiddenMarkovAccumulatorFactory

HiddenMarkovModelEstimator

alias of HiddenMarkovEstimator

HiddenMarkovModelDataEncoder

alias of HiddenMarkovDataEncoder

class HiddenMarkovFisherView(dist)[source]

Bases: FixedFisherView

Observed Fisher view for HMMs via forward-backward statistics.

The per-observation vectors are posterior-expected complete-data sufficient statistics: initial-state counts, transition counts, per-state emission statistics, optional length statistics, and state occupancies when the model accumulator exposes them. For finite enumerable HMMs, the full model Fisher is the exact observed covariance of these statistics under the model distribution. For continuous or otherwise non-enumerable HMMs, diagonal model moments remain available; use observed_fisher_information() for empirical full covariance on data.

Parameters:

dist (Any)

structured_statistics(x, estimate=None, weight=1.0)[source]

Return Baum-Welch sufficient statistics for one weighted sequence.

Parameters:
Return type:

Any

fisher_information(stats=None, diagonal=False, ridge=1.0e-8, **kwargs)[source]

Return HMM Fisher information, using exact finite-support moments when available.

Parameters:
Return type:

ndarray

fisher_vectors(stats=None, metric='diagonal', center=None, fisher=None, ridge=1.0e-8, **kwargs)[source]

Return Fisher-normalized HMM sufficient statistics.

Full Fisher normalization is available when the model has finite enumerable support or when an explicit Fisher matrix is supplied. Otherwise callers should use the diagonal metric or observed Fisher routes.

Parameters:
Return type:

ndarray

class HierarchicalNormalDistribution(mu, tau, sigma, name=None, keys=None, fit_diagnostics=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Two-level normal hierarchy over groups: y[g,i] ~ N(theta[g], sigma^2), theta[g] ~ N(mu, tau^2).

Each observation is one group (a sequence of values). mu is the population mean, tau the between-group sd and sigma the within-group sd. The latent group mean is marginalized out, so a group’s likelihood is the multivariate normal N(mu*1, sigma^2 I + tau^2 11^T) (computed in closed form). group_posterior / shrinkage give the partial-pooling estimates.

Parameters:

  • mu (float)

  • tau (float)

  • sigma (float)

  • name (str | None)

  • keys (str | None)

  • fit_diagnostics (HierarchicalNormalFitDiagnostics | None)

density(group)[source]

Return the marginal density of one observed group.

Return type:

float

log_density(group)[source]

Marginal log-likelihood of one group (latent group mean integrated out).

Return type:

float

seq_log_density(x)[source]

Return vectorized marginal log likelihoods for encoded groups.

Return type:

ndarray

group_posterior(ybar, n)[source]

Posterior (mean, sd) of a group’s true mean given its sample mean ybar and size n.

The shrinkage estimate mu + shrink*(ybar - mu) with shrink = tau^2/(tau^2 + sigma^2/n).

Parameters:
Return type:

tuple[float, float]

shrinkage(n)[source]

The shrinkage weight for a size-n group (0 = full pooling to mu, 1 = its own mean).

Parameters:

n (int)

Return type:

float

sampler(seed=None)[source]

Return a sampler for grouped observations from this hierarchy.

Parameters:

seed (int | None)

Return type:

HierarchicalNormalSampler

estimator(pseudo_count=None)[source]

Return an empirical-Bayes EM estimator for the hierarchy.

Parameters:

pseudo_count (float | None)

Return type:

HierarchicalNormalEstimator

dist_to_encoder()[source]

Return the data encoder used by this distribution.

Return type:

HierarchicalNormalDataEncoder

class HierarchicalNormalEstimator(max_iter=500, tol=1e-9, max_group_sizes=_DEFAULT_MAX_GROUP_SIZES, name=None, keys=None)[source]

Bases: ParameterEstimator

Weighted empirical-Bayes EM over size-bucketed group sufficient statistics.

Parameters:

  • max_iter (int)

  • tol (float)

  • max_group_sizes (int)

  • name (str | None)

  • keys (str | None)

accumulator_factory()[source]

Return an accumulator factory for bounded size-bucketed statistics.

Return type:

StatisticAccumulatorFactory

estimate(nobs, suff_stat)[source]

Fit mu, tau, and sigma with a weighted EM receipt.

Return type:

HierarchicalNormalDistribution

class HierarchicalNormalFitDiagnostics(converged, identifiable, iterations, termination_reason, objective_trace, final_parameter_delta, n_group_sizes, total_group_weight, total_observation_weight)[source]

Bases: object

Machine-readable receipt for the empirical-Bayes EM fit.

Parameters:

  • converged (bool)

  • identifiable (bool)

  • iterations (int)

  • termination_reason (str)

  • objective_trace (tuple[float, ...])

  • final_parameter_delta (float | None)

  • n_group_sizes (int)

  • total_group_weight (float)

  • total_observation_weight (float)

class HierarchicalMixtureDistribution(topics, mixture_weights, topic_weights, len_dist=NullDistribution(), name=None, keys=(None, None))[source]

Bases: SequenceEncodableProbabilityDistribution

Outer mixture over sequence mixtures with shared topic distributions.

Data type: Sequence[T], where T is the data type of the topic distributions.

Parameters:
compute_capabilities()[source]

Declare generated-compute support inherited from topics and length model.

compute_declaration()[source]

Return the generated-compute declaration for the hierarchical mixture.

density(x)[source]

Evaluate the density of an observation from hierarchical mixture distribution.

Parameters:

x (Sequence[T]) – A sequence of type data type T’s.

Returns:

Density evaluated at x.

Return type:

float

log_density(x)[source]

Evaluate the log density of an observation from hierarchical mixture distribution.

Note: Observation is a sequence.

Parameters:

x (Sequence[T]) – A sequence of type data type T’s.

Returns:

Log-density evaluated at x.

Return type:

float

posterior(x)[source]

Compute the posterior over the mixture components for the outer-mixture at observed value x.

Parameters:

x (Sequence[T]) – An observed sequence of data type T.

Returns:

Numpy array of length ‘num_mixtures’.

Return type:

ndarray

component_log_density(x)[source]

Evaluate the component-wise log-density for an observation from a hierarchical mixture model.

Parameters:

x (Sequence[T]) – An observation from a hierarchical mixture model.

Returns:

Numpy array length of ‘num_mixtures’.

Return type:

ndarray

to_mixture()[source]

Return the equivalent flat mixture of sequence distributions.

Return type:

MixtureDistribution

seq_component_log_density(x)[source]

Vectorized evaluation of the outer-mixture component-wise log-density for an encoded sequence x.

This returns a numpy array with shape (rv[0], ‘num_mixtures’).

Note: This density is a Mixture of Sequence of Mixture, so the data must be bin-counted as last step in code.

Encoded sequence ‘x’ is a Tuple of length 5 containing:

x[0] (int): Number of independent observations. x[1] (ndarray[int]): Observation sequence index for each value in flattened x. x[2] (ndarray[int]): Length of each observation in x. x[3] (E): Encoded sequence of flattened observed values (has type E). x[4] (Optional[E2]): Encoded sequence of lengths (has type E2).

Parameters:

x (tuple[int, ndarray, ndarray, E1, E2 | None]) – Encoded sequence of iid hierarchical mixture model observations.

Returns:

Numpy array of dimensions ‘rv[0]’ by ‘num_mixtures’, containing the log-density for each component of the

outer mixture.

Return type:

ndarray

seq_log_density(x)[source]

Vectorized evaluation of the log-density for an encoded sequence of observations in x.

Encoded sequence ‘x’ is a Tuple of length 5 containing:

x[0] (int): Number of independent observations. x[1] (ndarray[int]): Observation sequence index for each value in flattened x. x[2] (ndarray[int]): Length of each observation in x. x[3] (E): Encoded sequence of flattened observed values (has type E). x[4] (Optional[E2]): Encoded sequence of lengths (has type E2).

Parameters:

x (tuple[int, ndarray, ndarray, E1, E2 | None]) – Encoded sequence of observations of hierarchical mixture model.

Returns:

Log-density evaluated at each observation in the encoded sequence x.

Return type:

ndarray

backend_seq_component_log_density(x, engine)[source]

Engine-neutral outer-component log densities for hierarchical-mixture encoded sequences.

Parameters:
Return type:

Any

backend_seq_log_density(x, engine)[source]

Engine-neutral hierarchical-mixture log-density for encoded sequences.

Parameters:
Return type:

Any

seq_posterior(x)[source]

Vectorized evaluation of the posterior over each outer-mixture component for an encoded sequence x.

Encoded sequence ‘x’ is a Tuple of length 5 containing:

x[0] (int): Number of independent observations. x[1] (ndarray[int]): Observation sequence index for each value in flattened x. x[2] (ndarray[int]): Length of each observation in x. x[3] (E): Encoded sequence of flattened observed values (has type E). x[4] (Optional[E2]): Encoded sequence of lengths (has type E2).

Parameters:

x (tuple[int, ndarray, ndarray, E1, E2 | None]) – See above for details.

Returns:

Numpy array of dimension (x[0], ‘num_mixtures’).

Return type:

ndarray

to_fisher(**kwargs)[source]

Reuse the equivalent flat mixture’s Fisher view.

density_semantics()[source]

Return exact-or-approximate density semantics joined from child models.

sampler(seed=None)[source]

Return a sampler for this hierarchical mixture.

Parameters:

seed (int | None)

Return type:

HierarchicalMixtureSampler

estimator(pseudo_count=None)[source]

Create an estimator initialized from this hierarchical mixture.

Parameters:

pseudo_count (Optional[float]) – Prior mass used to smooth mixture and topic weights during estimation.

Returns:

Estimator configured with matching topic and length estimators.

Return type:

HierarchicalMixtureEstimator

dist_to_encoder()[source]

Return a data encoder for iid hierarchical-mixture observations.

Return type:

HierarchicalMixtureDataEncoder

enumerator()[source]

Returns a HierarchicalMixtureEnumerator iterating sequences in descending probability order.

Return type:

HierarchicalMixtureEnumerator

class HierarchicalMixtureEnumerator(dist)[source]

Bases: DistributionEnumerator

Enumerates the support of a HierarchicalMixtureDistribution in descending probability order.

Parameters:

dist (HierarchicalMixtureDistribution)

class HierarchicalMixtureSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for sequences from a hierarchical mixture distribution.

Parameters:

  • dist (HierarchicalMixtureDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Return samples from the underlying mixture sampler.

Parameters:
Return type:

Sequence[Any] | Any

HierarchicalMixtureAccumulator

alias of HierarchicalMixtureEstimatorAccumulator

HierarchicalMixtureAccumulatorFactory

alias of HierarchicalMixtureEstimatorAccumulatorFactory

class HierarchicalMixtureEstimator(estimators, num_mixtures, len_estimator=NullEstimator(), len_dist=None, suff_stat=None, pseudo_count=None, name=None, keys=(None, None))[source]

Bases: ParameterEstimator

Estimate hierarchical-mixture distributions from sufficient statistics.

Parameters:

  • estimators (Sequence[ParameterEstimator])

  • num_mixtures (int)

  • len_estimator (ParameterEstimator | None)

  • len_dist (SequenceEncodableProbabilityDistribution | None)

  • suff_stat (ndarray | None)

  • pseudo_count (float | None)

  • name (str | None)

  • keys (tuple[str | None, str | None] | None)

accumulator_factory()[source]

Return an accumulator factory configured from this estimator.

Return type:

HierarchicalMixtureEstimatorAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate a hierarchical mixture from aggregated sufficient statistics.

suff_stat is a four-item tuple containing:

suff_stat[0] (ndarray[float]): Aggregated token-level component counts with shape
    (num_mixtures, num_topics). Row-normalized into the topic weights taus.
suff_stat[1] (ndarray[float]): Aggregated document-level outer-posterior sums with length
    num_mixtures. Normalized into the outer weights w -- the EM maximizer for w
    is the document-level posterior sum, NOT the (length-weighted) token-level row sums.
suff_stat[2] (Tuple[SS1,...]): One sufficient-statistic value per topic.
suff_stat[3] (Optional[SS2]): Sufficient statistics for the length estimator.
Parameters:

  • nobs (Optional[float]) – Number of observations used in accumulation of ‘suff_stat’.

  • suff_stat (tuple[ndarray, ndarray, SS1, SS2 | None]) – See above for details.

Returns:

Estimated distribution.

Return type:

HierarchicalMixtureDistribution

class HierarchicalMixtureDataEncoder(topic_encoder, len_encoder)[source]

Bases: DataSequenceEncoder

Encode iid hierarchical-mixture observations for vectorized scoring.

Parameters:

  • topic_encoder (DataSequenceEncoder)

  • len_encoder (DataSequenceEncoder)

row_count(x)[source]

Return the explicit number of independent sequences in the payload.

Parameters:

x (tuple[int, ndarray, ndarray, Any, Any | None])

Return type:

int

seq_encode(x)[source]

Encode a sequence of iid observations from a hierarchical mixture model.

Returns ‘rv’ as a Tuple of length 5 containing:

rv[0] (int): Number of independent observations. rv[1] (ndarray[int]): Observation sequence index for each value in flattened x. rv[2] (ndarray[int]): Length of each observation in x. rv[3] (E): Encoded sequence of flattened observed values (has type E). rv[4] (Optional[E2]): Encoded sequence of lengths (has type E2).

Parameters:

x (Sequence[Sequence[T]]) – Sequence of hierarchical mixture model observations.

Returns:

See above.

Return type:

tuple[int, ndarray, ndarray, Any, Any | None]

determinize_quantized_terminal(dist, max_states=1 << 16)[source]

Determinize a terminal-value quantized HMM into a DeterminizedSequenceDistribution.

Raises EnumerationError if the HMM has no terminal_values, or if the belief expansion exceeds max_states (the twins property fails – not finitely determinizable).

Parameters:

max_states (int)

determinize_terminal_hmm(dist, max_states=1 << 16, max_denominator=10**9)[source]

Determinize a GENERAL terminal-value HMM (any transitions/initial, finite-discrete emissions).

Belief equality must be decidable, so each row of the model’s float probabilities is first rationalized (Fraction(p).limit_denominator(max_denominator)) and then normalized exactly. Determinization is exact on that stochastic rational model, which approximates the source to the requested rationalization precision. Requires terminal_values and enumerable finite emissions. Raises EnumerationError if not finitely determinizable within max_states or if an emission support is not finite/enumerable.

Parameters:

  • max_states (int)

  • max_denominator (int)

class DeterminizedSequenceDistribution(trans, accept, name=None)[source]

Bases: SequenceEncodableProbabilityDistribution

A deterministic weighted machine over terminal-ended sequences (one path per sequence).

trans[q][x] = (log_weight, next_state) for a non-terminal symbol; accept[q][x] = log_weight for a terminal symbol that completes the sequence. log_density(x) is the unique path’s summed log-weight (== the original HMM’s exact marginal); enumeration yields exact, duplicate-free n-best-strings.

Parameters:

name (str | None)

density(x)[source]

Return the probability of a terminal-ended sequence.

Return type:

float

log_density(x)[source]

Return the unique deterministic path log-weight for a sequence.

Return type:

float

seq_log_density(x)[source]

Return log-densities for a batch of encoded sequences.

Return type:

ndarray

dist_to_encoder()[source]

Return the passthrough sequence encoder for determinized machines.

Return type:

DeterminizedDataEncoder

enumerator()[source]

Return an exact descending-probability enumerator.

Return type:

DeterminizedEnumerator

quantized_count_index(quantizer, max_fine_bucket)[source]

Structural count index over the deterministic machine -> sub-linear exact seek/rank.

The machine is deterministic (one path per sequence), so a forward reach DP over states counts each sequence exactly once at its (floored) cost: reach_by_len[t][q] is the count histogram of length-t non-terminal prefixes reaching state q; an accept edge then completes a length-(t+1) sequence. Unranking walks the reach DP backward through the (precomputed) incoming edges. Exact up to quantization tie granularity; no path over-count (the determinization already removed it).

Parameters:

max_fine_bucket (int)

sampler(seed=None)[source]

Return a sampler that walks the deterministic weighted machine.

Parameters:

seed (int | None)

Return type:

DeterminizedSampler

estimator(pseudo_count=None)[source]

Raise because determinized machines are derived views, not fitted families.

Parameters:

pseudo_count (float | None)

class DeterminizedSampler(dist, seed=None)[source]

Bases: object

Generative sampler: at each state the accept+transition edge weights are the conditional next-symbol distribution (they sum to 1), so walk it until a terminal (accept) edge is taken.

Parameters:

  • dist (DeterminizedSequenceDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw one terminal-ended sequence or a list of independent sequences.

Parameters:
class DeterminizedDataEncoder[source]

Bases: DataSequenceEncoder

Passthrough encoder for determinized sequence observations.

seq_encode(x)[source]

Return sequence observations as a list without transforming symbols.

class DeterminizedEnumerator(dist)[source]

Bases: DistributionEnumerator

Exact descending-probability enumeration of sequences over the deterministic machine (n-best strings). Best-first with an admissible per-state best-completion bound (Viterbi over the machine).

Parameters:

dist (DeterminizedSequenceDistribution)

seek(index)[source]

Sub-linear exact seek over the deterministic machine (no prefix enumeration).

Deepens the structural count index until index is covered or the support is provably exhausted (the count-DP’s reliable truncated flag), then unranks within the located fine bucket. Robust to gaps in the cost spectrum, unlike the shared total-growth deepening heuristic.

Parameters:

index (int)

class IndianBuffetProcessFisherView(dist)[source]

Bases: FixedFisherView

Fixed Fisher view whose statistics are the binary feature indicators.

Parameters:

dist (Any)

class IndianBuffetProcessDistribution(num_features, alpha=1.0, beta_params=None, feature_probs=None, min_prob=1.0e-128, name=None, keys=None, data_format='auto')[source]

Bases: SequenceEncodableProbabilityDistribution

Finite-truncated Indian buffet process over binary feature rows.

Parameters:

  • num_features (int) – Truncation level K.

  • alpha (float) – IBP concentration parameter.

  • beta_params (Sequence[Sequence[float]] | ndarray | None) – Optional (K, 2) variational Beta parameters for q(pi_k). If omitted, the prior Beta(alpha / K, 1) is used.

  • feature_probs (Sequence[float] | ndarray | None) – Optional plug-in feature probabilities. When supplied without beta_params, a lightweight Beta posterior with matching mean is created for expected-log-density calls.

  • min_prob (float) – Minimum plug-in probability used when feature_probs are given.

  • name (str | None) – Optional distribution name.

  • keys (str | None) – Optional key for tying sufficient statistics.

  • data_format (str) – ‘dense’, ‘sparse’, or ‘auto’ input interpretation.

classmethod compute_capabilities()[source]

Declare generic table-kernel capabilities for finite IBP rows.

classmethod compute_declaration()[source]

Return the structured declaration for finite IBP parameters and statistics.

density(x)[source]

Return the probability density or mass at a single observation.

Parameters:

x (Any)

Return type:

float

density_semantics()[source]

Return that plug-in densities use posterior mean feature probabilities.

log_density(x)[source]

Plug-in log-density of one feature row using E_q[pi_k].

Parameters:

x (Any)

Return type:

float

expected_log_density(x)[source]

VB expected log-density E_q[log p(z | pi)] for one observed row.

Parameters:

x (Any)

Return type:

float

seq_log_density(x)[source]

Return vectorized log-density values for sequence-encoded observations.

Parameters:

x (ndarray)

Return type:

ndarray

backend_seq_log_density(x, engine)[source]

Engine-neutral vectorized plug-in log-density for encoded feature rows.

Parameters:
Return type:

Any

classmethod backend_stacked_params(dists, engine)[source]

Return stacked finite-IBP parameters for equal-feature mixtures.

Parameters:

  • dists (Sequence[IndianBuffetProcessDistribution])

  • engine (Any)

Return type:

dict[str, Any]

classmethod backend_stacked_log_density(x, params, engine)[source]

Return an (n, k) matrix of finite-IBP component log densities.

Parameters:
Return type:

Any

classmethod backend_stacked_sufficient_statistics(x, weights, params, engine)[source]

Return component-stacked legacy (feature_counts, total_count, alpha) statistics.

Parameters:
Return type:

tuple[Any, Any, Any]

seq_expected_log_density(x)[source]

Return vectorized expected log-density values for encoded observations.

Parameters:

x (ndarray)

Return type:

ndarray

seq_local_elbo(x)[source]

Per-row VB contribution; rows are observed, so there is no local entropy.

Parameters:

x (ndarray)

Return type:

ndarray

enumerator()[source]

Enumerate feature rows in descending probability order.

The plug-in row density factorizes over features – log p(z) = sum_k [z_k log pi_k + (1-z_k) log(1-pi_k)] with pi_k = E_q[feature_probs] – so the truncated IBP row is a product of independent Bernoulli features and enumerates by best-first over the per-feature supports (the same structure as the Erdos-Renyi graph). Rows are emitted in the configured data_format (a dense 0/1 list, or a sorted list of active feature indices when sparse), each carrying its exact log_density.

Return type:

DistributionEnumerator

to_fisher(**kwargs)[source]

Return this distribution’s own Fisher view.

sampler(seed=None)[source]

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

IndianBuffetProcessSampler

estimator(pseudo_count=None)[source]

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

IndianBuffetProcessEstimator

dist_to_encoder()[source]

Return the data encoder used by this distribution for vectorized methods.

Return type:

IndianBuffetProcessDataEncoder

class IndianBuffetProcessEnumerator(dist)[source]

Bases: DistributionEnumerator

Enumerate finite IBP feature rows in descending plug-in probability order.

Parameters:

dist (IndianBuffetProcessDistribution)

class IndianBuffetProcessSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for finite-truncated IBP feature rows.

Parameters:

  • dist (IndianBuffetProcessDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw one feature row or size iid feature rows.

Parameters:
Return type:

list[int] | list[list[int]]

class IndianBuffetProcessAccumulator(num_features, alpha=1.0, keys=None, data_format='auto')[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulates weighted feature-use counts for the IBP VB update.

Parameters:

  • num_features (int)

  • alpha (float)

  • keys (str | None)

  • data_format (str)

update(x, weight, estimate)[source]

Update weighted feature-use counts from one row.

Parameters:

  • x (Any)

  • weight (float)

  • estimate (IndianBuffetProcessDistribution | None)

Return type:

None

initialize(x, weight, rng)[source]

Initialize feature-use counts from one row.

Parameters:
Return type:

None

seq_update(x, weights, estimate)[source]

Update weighted feature-use counts from encoded rows.

Parameters:

  • x (ndarray)

  • weights (ndarray)

  • estimate (IndianBuffetProcessDistribution | None)

Return type:

None

seq_update_engine(x, weights, estimate, engine)[source]

Engine-resident accumulation of weighted feature-use counts (numpy or torch).

The weighted feature counts are reduced via a weight-vector / binary-matrix product on the active engine; the alpha metadata is host bookkeeping. Matches seq_update.

Parameters:

  • x (ndarray)

  • weights (Any)

  • estimate (IndianBuffetProcessDistribution | None)

  • engine (Any)

Return type:

None

seq_initialize(x, weights, rng)[source]

Initialize feature-use counts from encoded rows.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge feature-use counts, total count, and alpha metadata.

Parameters:

suff_stat (tuple[ndarray, float, float | None])

Return type:

IndianBuffetProcessAccumulator

value()[source]

Return feature-use counts, total count, and alpha metadata.

Return type:

tuple[ndarray, float, float | None]

from_value(x)[source]

Restore feature-use counts, total count, and alpha metadata.

Parameters:

x (tuple[ndarray, float, float | None])

Return type:

IndianBuffetProcessAccumulator

scale(c)[source]

Scale additive counts while preserving the IBP concentration metadata.

Parameters:

c (float)

Return type:

IndianBuffetProcessAccumulator

acc_to_encoder()[source]

Return the encoder compatible with finite IBP sufficient statistics.

Return type:

IndianBuffetProcessDataEncoder

class IndianBuffetProcessAccumulatorFactory(num_features, alpha=1.0, keys=None, data_format='auto')[source]

Bases: StatisticAccumulatorFactory

Factory for IBP accumulators.

Parameters:

  • num_features (int)

  • alpha (float)

  • keys (str | None)

  • data_format (str)

make()[source]

Create an empty finite IBP accumulator.

Return type:

IndianBuffetProcessAccumulator

class IndianBuffetProcessEstimator(num_features, alpha=1.0, pseudo_count=None, suff_stat=None, estimate_alpha=True, min_alpha=1.0e-12, max_alpha=1.0e12, min_prob=1.0e-128, name=None, keys=None, data_format='auto')[source]

Bases: ParameterEstimator

Variational Bayes estimator for the finite-truncated IBP.

pseudo_count follows the convention used by other mixle.stats Bernoulli estimators: if suff_stat is supplied, it is treated as a prior probability vector and re-weighted by pseudo_count; otherwise pseudo_count is centered at the IBP prior mean alpha / (alpha + K).

Parameters:
accumulator_factory()[source]

Return a factory for finite IBP sufficient-statistic accumulators.

Return type:

IndianBuffetProcessAccumulatorFactory

estimate(nobs, suff_stat)[source]

Update the variational Beta posterior and optional concentration parameter.

Parameters:
Return type:

IndianBuffetProcessDistribution

model_log_density(model)[source]

Global VB term E_q[log p(pi | alpha)] + H[q(pi)].

Parameters:

model (IndianBuffetProcessDistribution)

Return type:

float

class IndianBuffetProcessDataEncoder(num_features, data_format='auto')[source]

Bases: DataSequenceEncoder

Encode dense or sparse IBP rows as a dense boolean matrix.

Parameters:

  • num_features (int)

  • data_format (str)

seq_encode(x)[source]

Validate and encode dense or sparse feature rows as a boolean matrix.

Parameters:

x (Sequence[Any] | ndarray)

Return type:

ndarray

class IntegerHiddenAssociationStatistics(initial_counts, weight_counts, state_counts, previous, size, observation_mass, schema_version=1)[source]

Bases: Sequence[Any]

Mass-aware statistics with backward five-slot iteration.

Parameters:
class IntegerHiddenAssociationDistribution(state_prob_mat, cond_weights, alpha=0.0, prev_dist=NullDistribution(), len_dist=NullDistribution(), name=None, keys=(None, None), use_numba=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Integer hidden association model: words of a second set are emitted through hidden states conditioned on words of a first set.

Parameters:
compute_capabilities()[source]

Return backend capability metadata for this concrete integer association model.

compute_declaration()[source]

Return the generated-compute declaration for integer hidden association.

log_density(x)[source]

Log-density of the integer hidden association model at observation x.

For each emitted word in x[1], marginalizes over the given words in x[0] (weighted by count) and the hidden states, mixing with a uniform density with probability alpha. Adds the log-density of x[0] under prev_dist and of the total emission count under len_dist.

Parameters:

x (Tuple[List[Tuple[int, float]], List[Tuple[int, float]]]) – Grouped-count observation ([(S1 word, count)], [(S2 word, count)]).

Returns:

Log-density at observation x.

Return type:

float

seq_log_density(x)[source]

Vectorized evaluation of log-density at sequence encoded input x.

Parameters:

x (E) – Sequence encoded observations from IntegerHiddenAssociationDataEncoder.seq_encode(). Uses the numba kernel when the encoding was produced with use_numba=True.

Returns:

Numpy array of log-density values, one per encoded observation.

Return type:

ndarray

backend_seq_log_density(x, engine)[source]

Evaluate encoded log-densities using a backend-neutral compute engine.

Parameters:
Return type:

Any

conditional_word_log_probs(s1)[source]

Log of the per-emission word distribution q(.|S1) for a given S1 bag, or None if empty.

q(w|S1) = (1-alpha) * sum_u (c_u/n1) * sum_s cond_weights[u,s] * state_prob_mat[s,w] + alpha/W – the smoothed mixture the model uses to score each emitted word. Returns None for an empty S1 (n1 = 0), whose conditional is degenerate (the model’s own density is undefined there).

Parameters:

s1 (list[tuple[int, float]])

Return type:

ndarray | None

enumerator()[source]

Enumerate (S1, S2) observations in descending probability order.

The model factors as prev_dist(S1) * [prod_w q(w|S1)^{c_w}] * P_len(n2): the emitted bag S2 is a trial-count multinomial whose word distribution q(.|S1) depends on the given bag S1. Enumeration is a conditional product – the outer stream enumerates S1 from prev_dist and, for each S1, the inner stream enumerates S2 by the multinomial bag search under len_dist, merged by descending total score with prev_dist(S1) as the outer frontier bound. Requires an enumerable, non-null prev_dist so the S1 support is defined.

Return type:

DistributionEnumerator

sampler(seed=None)[source]

Return a sampler for iid integer hidden-association observations.

Requires non-null prev_dist and len_dist.

Parameters:

seed (int | None) – Optional random seed.

Returns:

A configured integer hidden-association sampler.

Return type:

IntegerHiddenAssociationSampler

estimator(pseudo_count=None)[source]

Return an estimator with matching dimensions and child estimators.

Parameters:

pseudo_count (float | None) – Unused; accepted for protocol consistency.

Returns:

A configured integer hidden-association estimator.

Return type:

IntegerHiddenAssociationEstimator

dist_to_encoder()[source]

Return an encoder for integer hidden-association observations.

Return type:

IntegerHiddenAssociationDataEncoder

class IntegerHiddenAssociationEnumerator(dist)[source]

Bases: DistributionEnumerator

Enumerate integer hidden-association support pairs in best-first order.

Parameters:

dist (IntegerHiddenAssociationDistribution)

class IntegerHiddenAssociationSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for grouped-count word-set pairs from an integer hidden-association distribution.

Parameters:

  • dist (IntegerHiddenAssociationDistribution)

  • seed (int | None)

sample_given(x)[source]

Draw an emitted grouped-count word set conditioned on the given word set x.

Parameters:

x (List[Tuple[int, float]]) – Given word set as (word, count) pairs.

Returns:

List of (emitted word, count) pairs.

Return type:

list[tuple[int, float]]

sample(size=None, *, batched=True)[source]

Draw iid grouped-count observations from the integer hidden association model.

Parameters:

  • size (Optional[int]) – Number of observations to draw. If None, a single observation is returned.

  • batched (bool)

Returns:

A ([(S1 word, count)], [(S2 word, count)]) tuple if size is None, else a list of such tuples of length size.

Return type:

Sequence[list[tuple[int, float]]] | list[tuple[int, float]]

class IntegerHiddenAssociationAccumulator(num_vals1, num_vals2, num_states, prev_acc=NullAccumulator(), size_acc=NullAccumulator(), use_numba=None, keys=(None, None))[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulator for hidden-state and emission counts from observed word-set pairs.

Parameters:

  • num_vals1 (int)

  • num_vals2 (int)

  • num_states (int)

  • prev_acc (SequenceEncodableStatisticAccumulator | None)

  • size_acc (SequenceEncodableStatisticAccumulator | None)

  • use_numba (bool | None)

  • keys (tuple[str | None, str | None] | None)

update(x, weight, estimate)[source]

Update sufficient statistics with posterior word/state assignments for the observation.

Parameters:

  • x (Tuple[List[Tuple[int, float]], List[Tuple[int, float]]]) – Grouped-count observation ([(S1 word, count)], [(S2 word, count)]).

  • weight (float) – Weight for the observation.

  • estimate (IntegerHiddenAssociationDistribution) – Previous estimate used to compute posteriors.

Return type:

None

initialize(x, weight, rng)[source]

Initialize sufficient statistics with random (Dirichlet) state assignments.

Parameters:

  • x (Tuple[List[Tuple[int, float]], List[Tuple[int, float]]]) – Grouped-count observation ([(S1 word, count)], [(S2 word, count)]).

  • weight (float) – Weight for the observation.

  • rng (np.random.RandomState) – Random number generator for the random assignments.

Return type:

None

seq_initialize(x, weights, rng)[source]

Vectorized initialization of sufficient statistics from sequence encoded observations.

Parameters:

  • x (E) – Sequence encoded observations from IntegerHiddenAssociationDataEncoder.seq_encode().

  • weights (np.ndarray) – Weights, one per encoded observation.

  • rng (np.random.RandomState) – Random number generator for the random assignments.

Return type:

None

seq_update(x, weights, estimate)[source]

Vectorized update of sufficient statistics from sequence encoded observations.

Parameters:

  • x (E) – Sequence encoded observations from IntegerHiddenAssociationDataEncoder.seq_encode(). Uses the numba kernel when the encoding was produced with use_numba=True.

  • weights (np.ndarray) – Weights, one per encoded observation.

  • estimate (IntegerHiddenAssociationDistribution) – Previous estimate used to compute posteriors.

Return type:

None

seq_update_engine(x, weights, estimate, engine)[source]

Engine-resident E-step for the pure (non-numba) blocked encoding.

Mirrors the numpy branch of seq_update: for each observation the (given word x state x emitted word) responsibility tensor is built and normalized with the alpha smoothing on the active engine (numpy or torch), and the initial/weight/state counts are scattered into engine-resident accumulators via index_add. Only the per-observation orchestration runs in Python; all tensor arithmetic and accumulation are on the engine.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge sufficient statistics of suff_stat into this accumulator.

Parameters:

suff_stat (Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[SS1], Optional[SS2]]) – Init counts, weight counts, state counts, prev suff stats, and size suff stats.

Returns:

This IntegerHiddenAssociationAccumulator.

Return type:

IntegerHiddenAssociationAccumulator

value()[source]

Returns the sufficient statistics: (init counts, weight counts, state counts, prev, size).

Return type:

IntegerHiddenAssociationStatistics

from_value(x)[source]

Set the sufficient statistics of this accumulator from x.

Parameters:

x (Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[SS1], Optional[SS2]]) – Init counts, weight counts, state counts, prev suff stats, and size suff stats.

Returns:

This IntegerHiddenAssociationAccumulator.

Return type:

IntegerHiddenAssociationAccumulator

scale(c)[source]

Scale linear association counts and delegate child accumulators.

Parameters:

c (float)

Return type:

IntegerHiddenAssociationAccumulator

key_merge(stats_dict)[source]

Merge this accumulator’s weight and state counts into stats_dict under their keys, if keyed.

Parameters:

stats_dict (Dict[str, Any]) – Maps keys to merged sufficient statistics.

Return type:

None

key_replace(stats_dict)[source]

Replace this accumulator’s weight and state counts with the keyed statistics in stats_dict, if keyed.

Parameters:

stats_dict (Dict[str, Any]) – Maps keys to merged sufficient statistics.

Return type:

None

acc_to_encoder()[source]

Return an encoder compatible with integer hidden-association observations.

Return type:

DataSequenceEncoder

class IntegerHiddenAssociationAccumulatorFactory(num_vals1, num_vals2, num_states, prev_factory=NullAccumulatorFactory(), len_factory=NullAccumulatorFactory(), use_numba=None, keys=(None, None))[source]

Bases: StatisticAccumulatorFactory

Factory for integer hidden-association accumulators.

Parameters:

  • num_vals1 (int)

  • num_vals2 (int)

  • num_states (int)

  • prev_factory (StatisticAccumulatorFactory | None)

  • len_factory (StatisticAccumulatorFactory | None)

  • use_numba (bool | None)

  • keys (tuple[str | None, str | None])

make()[source]

Return a fresh integer hidden-association accumulator.

Return type:

IntegerHiddenAssociationAccumulator

class IntegerHiddenAssociationEstimator(num_vals, num_states, alpha=0.0, prev_estimator=NullEstimator(), len_estimator=NullEstimator(), suff_stat=None, pseudo_count=None, use_numba=None, name=None, keys=(None, None))[source]

Bases: ParameterEstimator

Estimator for integer hidden-association distributions.

Parameters:

  • num_vals (list[int] | tuple[int, int] | int)

  • num_states (int)

  • alpha (float)

  • prev_estimator (ParameterEstimator | None)

  • len_estimator (ParameterEstimator | None)

  • suff_stat (Any | None)

  • pseudo_count (float | None)

  • use_numba (bool | None)

  • name (str | None)

  • keys (tuple[str | None, str | None] | None)

accumulator_factory()[source]

Return an accumulator factory matching this estimator.

Return type:

IntegerHiddenAssociationAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate an IntegerHiddenAssociationDistribution from aggregated sufficient statistics.

Parameters:

  • nobs (Optional[float]) – Number of observations, passed to the prev and length estimators.

  • suff_stat (Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[SS1], Optional[SS2]]) – Init counts, weight counts, state counts, prev suff stats, and size suff stats.

Returns:

IntegerHiddenAssociationDistribution object.

Return type:

IntegerHiddenAssociationDistribution

class IntegerHiddenAssociationDataEncoder(num_vals1, num_vals2, prev_encoder, len_encoder, use_numba)[source]

Bases: DataSequenceEncoder

Encode grouped-count integer hidden-association observations.

Parameters:

  • num_vals1 (int)

  • num_vals2 (int)

  • prev_encoder (DataSequenceEncoder)

  • len_encoder (DataSequenceEncoder)

  • use_numba (bool)

seq_encode(x)[source]

Sequence encoding for integer hidden association observations.

If numba is not used see _seq_encode(). Else the following is returned a Tuple of the following form is returned None, ((s0, s1, x0, x1, c0, c1, w0), xv, nn) with,

s0 (np.ndarray): Numpy array of lengths for length of x[i][0] s1 (np.ndarray): Numpy array of lengths for length of x[i][1]. x0 (np.ndarray): Flattened numpy array of values from x[i][0]. x1 (np.ndarray): Flattened numpy array of values from x[i][1]. c0 (np.ndarray): Flattened numpy array of counts from x[i][0]. c1 (np.ndarray): Flattened numpy array of counts from x[i][1]. w0 (np.ndarray): Numpy array of sum of counts for each x[i][0]. xv (E1): Sequence encoded flattened values of x[i][0]. nn (E2): Sequence encoded values of lengths (counts).

Parameters:

x (Sequence[tuple[list[tuple[int, float]], list[tuple[int, float]]]]) – Sequence of iid integer hidden association observations.

Returns:

See above.

Return type:

tuple[tuple[list[tuple[ndarray, …]], Any | None, Any | None], None] | tuple[None, tuple[ndarray, …], Any | None, Any | None]

row_count(x)[source]

Return the number of validated association rows in an encoded payload.

Parameters:

x (Any)

Return type:

int

class IntegerPLSIEncodedData(value_ids, counts, entry_doc_ids, row_ids, totals, doc_ids, num_vals, num_docs)[source]

Bases: object

Owned, validated flattened representation of canonical integer PLSI bags.

Parameters:
class IntegerProbabilisticLatentSemanticIndexingDistribution(state_word_mat, doc_state_mat, doc_vec, len_dist=NullDistribution(), name=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Integer-valued probabilistic latent semantic indexing distribution.

Parameters:
compute_capabilities()[source]

Return backend capability metadata for this concrete PLSI instance.

compute_declaration()[source]

Return the symbolic distribution declaration for code generation and PPL introspection.

density(x)[source]

Evaluate the density of PLSI model for an observation x.

See log_density() for details on the density evaluation.

Parameters:

x (Tuple[int, Sequence[Tuple[int, float]]]) – Single observation of integer PLSI.

Returns:

Density evaluated at observed value x.

Return type:

float

log_density(x)[source]

Evaluate the log-density of PLSI model for an observation of x.

Consider an Integer PLSI model for a corpus of documents with S states, V word values, and D documents ids (authors).

Let x (Tuple[int, Sequence[Tuple[int, float]]]) be an observation from a PLSI model, consisting of x = (d, [(v_0, c_0), (v_1, c_1), …, (v_{k-1}, c_{k-1})]), where the ‘d’ is some document d_id in the corpus and each tuple (v_i, c_i) corresponds to a value-count couple in the corpus. The log-likelihood is given by

log(p_mat(x)) = log(p_mat(d)) + sum_{j=0}^{k-1} c_k*log( sum_{s=0}^{S-1} p_mat(d|s)p_mat(s|v_k) ) + log(P_len(nn)),

where P_len(nn) is the density of the length distribution for ‘nn’ representing the total number of words in the document.

Parameters:

x (Tuple[int, Sequence[Tuple[int, float]]]) – (doc_id, [(value_id, count_for_value)]). See above for details.

Returns:

Log-density evaluated at a single observation x.

Return type:

float

component_log_density(x)[source]

Evaluate the log-density for each state in the PLSI.

Returns count*log(p_mat(W|S)) for each word-count pair in the document. Returned value is S by 1 where S is the number of components in the model.

Parameters:

x (Tuple[int, Sequence[Tuple[int, float]]]) – Single PLSI observation of form (doc_id, [(value_id, count_for_value)]).

Returns:

Numpy array of length S (num_states).

Return type:

ndarray

seq_log_density(x)[source]

Vectorized evaluation of the log-density for an encoded sequence of iid observation from a PLSI model.

See log_density() function for details on the log-likelihood.

The encoded sequence ‘x’ is a Tuple length 2. The first component contains data type Optional[T1] corresponding to the sequence encoding of the lengths. The second component is a Tuple of length 6 containing

xv (ndarray[int]): Numpy array of flattened word values.
xc (ndarray[float]): Numpy array of flattened counts for word values above.
xd (ndarray[int]): Document id for each word-count pair in the arrays above.
xi (ndarray[int]): Observed sequence index for each word-count pair in the arrays above.
xn (ndarray[float]): Numpy array of the total number of words in each document.
xm (ndarray[float]): Flattened array of document id's for the lengths above (len = len(x)).
Parameters:

x (tuple[T1 | None, tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray]]) – Encoded sequence of iid observations of PLSI model. See above for details.

Returns:

Numpy array of log-density evaluated at each observation in the encoded sequence.

Return type:

ndarray

backend_seq_log_density(x, engine)[source]

Evaluate encoded PLSI log densities using a backend-neutral compute engine.

Parameters:
Return type:

Any

seq_component_log_density(x)[source]
Vectorized evaluation of the component log-density for each observation in an encoded sequence of iid PLSI

observations.

See component_log_density() function for details on component log-likelihood evaluation.

The encoded sequence ‘x’ is a Tuple length 2. The first component contains data type Optional[T1] corresponding to the sequence encoding of the lengths. The second component is a Tuple of length 6 containing

xv (ndarray[int]): Numpy array of flattened word values.
xc (ndarray[float]): Numpy array of flattened counts for word values above.
xd (ndarray[int]): Document id for each word-count pair in the arrays above.
xi (ndarray[int]): Observed sequence index for each word-count pair in the arrays above.
xn (ndarray[float]): Numpy array of the total number of words in each document.
xm (ndarray[float]): Flattened array of document id's for the lengths above (len = len(x)).
Parameters:

x (tuple[T1 | None, tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray]]) – Encoded sequence of iid observations of PLSI model. See above for details.

Returns:

2-d numpy array containing N rows of num_state sized arrays.

Return type:

ndarray

enumerator()[source]

Enumerate PLSI observations (doc_id, bag) in descending probability order.

A PLSI observation factors as P(doc) * [prod_w q_d(w)^{c_w}] * P_len(n) where q_d is the per-document word distribution prob_mat @ state_mat[d] and n the total word count, so it is a document-labelled mixture of trial-count multinomials: for each document the bags enumerate by a multiset best-first search under a length frontier driven by len_dist (the real trial-count distribution), and the per-document streams are merged by descending score with the document log-probability as offset. Requires a modelled len_dist unless every per-document word distribution is sub-stochastic-free; an absent length distribution leaves the bag support infinite and is enumerated by the multinomial term alone.

Return type:

DistributionEnumerator

sampler(seed=None)[source]

Return a sampler for iid integer PLSI observations.

Parameters:

seed (int | None)

Return type:

IntegerProbabilisticLatentSemanticIndexingSampler

estimator(pseudo_count=None)[source]

Return an estimator initialized from this distribution’s dimensions.

Parameters:

pseudo_count (float | None) – Optional smoothing count for topic, word, and document counts.

Returns:

A configured integer PLSI estimator.

Return type:

IntegerProbabilisticLatentSemanticIndexingEstimator

dist_to_encoder()[source]

Return an encoder for integer PLSI observations.

Return type:

IntegerProbabilisticLatentSemanticIndexingDataEncoder

class IntegerProbabilisticLatentSemanticIndexingEnumerator(dist)[source]

Bases: DistributionEnumerator

Best-first enumerator for document-labelled integer PLSI bag observations.

Parameters:

dist (IntegerProbabilisticLatentSemanticIndexingDistribution)

class IntegerProbabilisticLatentSemanticIndexingSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for integer PLSI document ids, word-count bags, and document lengths.

Parameters:

  • dist (IntegerProbabilisticLatentSemanticIndexingDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Generate iid samples from PLSI model.

Parameters:

  • size (Optional[int]) – Number of samples to generate. Defaults to 0 if size is None.

  • batched (bool)

Returns:

Sequence of iid PLSI samples if size is not None, else a single sample from PLSI model.

Return type:

tuple[int, Sequence[tuple[int, float]]] | Sequence[tuple[int, Sequence[tuple[int, float]]]]

class IntegerProbabilisticLatentSemanticIndexingAccumulator(num_vals, num_states, num_docs, len_acc=NullAccumulator(), name=None, keys=(None, None, None))[source]

Bases: SequenceEncodableStatisticAccumulator

EM sufficient-statistic accumulator for integer PLSI word, state, document, and length terms.

Parameters:

  • num_vals (int)

  • num_states (int)

  • num_docs (int)

  • len_acc (SequenceEncodableStatisticAccumulator | None)

  • name (str | None)

  • keys (tuple[str | None, str | None, str | None] | None)

update(x, weight, estimate)[source]

Update sufficient statistics from one weighted sparse-bag observation.

Parameters:

  • x (tuple[int, Sequence[tuple[int, float]]]) – Integer PLSI observation as (doc_id, [(value_id, count), ...]).

  • weight (float) – Observation weight.

  • estimate (IntegerProbabilisticLatentSemanticIndexingDistribution) – Previous integer PLSI estimate.

Return type:

None

initialize(x, weight, rng)[source]

Initialize sufficient statistics from one weighted sparse-bag observation.

Parameters:
Return type:

None

seq_initialize(x, weights, rng)[source]

Vectorized initialization of sufficient statistics form an encoded sequence of observations in arg ‘x’.

The encoded sequence ‘x’ is a Tuple length 2. The first component contains data type Optional[T1] corresponding to the sequence encoding of the lengths. The second component is a Tuple of length 6 containing

xv (ndarray[int]): Numpy array of flattened word values.
xc (ndarray[float]): Numpy array of flattened counts for word values above.
xd (ndarray[int]): Document id for each word-count pair in the arrays above.
xi (ndarray[int]): Observed sequence index for each word-count pair in the arrays above.
xn (ndarray[float]): Numpy array of the total number of words in each document.
xm (ndarray[float]): Flattened array of document id's for the lengths above (len = len(x)).
Parameters:

  • x (tuple[T1 | None, tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray]]) – Encoded sequence of iid observations of PLSI. See above for details.

  • weights (ndarray) – Weights for observations in encoded sequence.

  • rng (RandomState) – Used to initialize member RandomState variables.

Returns:

None.

Return type:

None

seq_update(x, weights, estimate)[source]

Vectorized update of sufficient statistics for encoded sequence of iid observations in x.

The encoded sequence ‘x’ is a Tuple length 2. The first component contains data type Optional[T1] corresponding to the sequence encoding of the lengths. The second component is a Tuple of length 6 containing

xv (ndarray[int]): Numpy array of flattened word values.
xc (ndarray[float]): Numpy array of flattened counts for word values above.
xd (ndarray[int]): Document id for each word-count pair in the arrays above.
xi (ndarray[int]): Observed sequence index for each word-count pair in the arrays above.
xn (ndarray[float]): Numpy array of the total number of words in each document.
xm (ndarray[float]): Flattened array of document id's for the lengths above (len = len(x)).
Parameters:

  • x (tuple[T1 | None, tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray]]) – Encoded sequence of iid observations of PLSI model. See above for details.

  • weights (np.ndarray) – Numpy array of observation weights.

  • estimate (IntegerProbabilisticLatentSemanticIndexingDistribution) – Prior estimate of IntegerProbabilisticLatentSemanticIndexingDistribution object.

Returns:

None.

Return type:

None

seq_update_engine(x, weights, estimate, engine)[source]

Engine-resident E-step: the PLSI responsibility update (state-word x doc-state gather, per-pair normalization, and the word/doc segment sums) runs on the active engine, matching the host seq_update.

combine(suff_stat)[source]

Merge aggregated integer PLSI sufficient statistics into this accumulator.

The tuple is interpreted as (word_count, comp_count, doc_count, length_stats).

Parameters:

suff_stat (tuple[ndarray, ndarray, ndarray, SS1 | None]) – Aggregated integer PLSI sufficient statistics.

Returns:

This accumulator.

Return type:

IntegerProbabilisticLatentSemanticIndexingAccumulator

value()[source]

Return sufficient statistics as (word_count, comp_count, doc_count, length_stats).

Return type:

tuple[ndarray, ndarray, ndarray, Any | None]

from_value(x)[source]

Replace this accumulator’s sufficient statistics.

Parameters:

x (tuple[ndarray, ndarray, ndarray, SS1 | None]) – Aggregated sufficient statistics in value format.

Returns:

This accumulator.

Return type:

IntegerProbabilisticLatentSemanticIndexingAccumulator

scale(c)[source]

Scale linear latent counts and delegate document-length statistics.

Parameters:

c (float)

Return type:

IntegerProbabilisticLatentSemanticIndexingAccumulator

key_merge(stats_dict)[source]

Merge this accumulator into stats_dict under configured keys.

If wc_key is set, merge the state/word count variable. If sc_key is set, merge the doc/state count variable. If dc_key is set, merge the author count variable.

The length accumulator receives the same merge request.

Parameters:

stats_dict (dict[str, Any]) – Mapping from merge keys to sufficient statistics.

Return type:

None

key_replace(stats_dict)[source]

Replace sufficient statistics from matching keys in stats_dict.

If wc_key is set, set the state/word count variable to matching key in stats_dict. If sc_key is set, set the doc/state count variable to matching key in stats_dict. If dc_key is set, set the author count variable to matching key in stats_dict.

The length accumulator receives the same replace request.

Parameters:

stats_dict (dict[str, Any]) – Mapping from merge keys to sufficient statistics.

Return type:

None

acc_to_encoder()[source]

Return an encoder compatible with integer PLSI observations.

Return type:

IntegerProbabilisticLatentSemanticIndexingDataEncoder

class IntegerProbabilisticLatentSemanticIndexingAccumulatorFactory(num_vals, num_states, num_docs, len_factory=NullAccumulatorFactory(), keys=(None, None, None), name=None)[source]

Bases: StatisticAccumulatorFactory

Factory for integer PLSI EM sufficient-statistic accumulators.

Parameters:

  • num_vals (int)

  • num_states (int)

  • num_docs (int)

  • len_factory (StatisticAccumulatorFactory | None)

  • keys (tuple[str | None, str | None, str | None] | None)

  • name (str | None)

make()[source]

Return a fresh integer PLSI accumulator.

Return type:

IntegerProbabilisticLatentSemanticIndexingAccumulator

class IntegerProbabilisticLatentSemanticIndexingEstimator(num_vals, num_states, num_docs, len_estimator=NullEstimator(), pseudo_count=(None, None, None), suff_stat=(None, None, None), name=None, keys=(None, None, None))[source]

Bases: ParameterEstimator

Estimator for integer PLSI word/state/document probabilities and the optional length model.

Parameters:
accumulator_factory()[source]

Return an accumulator factory matching this estimator.

Return type:

IntegerProbabilisticLatentSemanticIndexingAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate an integer PLSI distribution from aggregated sufficient statistics.

Parameters:

  • nobs (float | None) – Optional observation count, accepted for the estimator interface.

  • suff_stat (tuple[ndarray, ndarray, ndarray, SS1 | None]) – Aggregated word, topic, document, and length statistics.

Returns:

A fitted integer PLSI distribution.

Return type:

IntegerProbabilisticLatentSemanticIndexingDistribution

class IntegerProbabilisticLatentSemanticIndexingDataEncoder(num_vals, num_docs, len_encoder=NullDataEncoder())[source]

Bases: DataSequenceEncoder

Encode integer PLSI observations into flattened sparse bag arrays and length features.

Parameters:

  • num_vals (int)

  • num_docs (int)

  • len_encoder (DataSequenceEncoder | None)

seq_encode(x)[source]

Encode iid PLSI observations for vectorized seq_* methods.

Input arg ‘x’ is a sequence of iid PLSI observations having form

x = [ (doc_id, [(value, count),…]),… ].

The return value has two entries. The first contains the optional length encoding. The second contains:

xv (ndarray[int]): Numpy array of flattened word values. xc (ndarray[float]): Numpy array of flattened counts for word values above. xd (ndarray[int]): Document d_id for each word-count pair in the arrays above. xi (ndarray[int]): Observed sequence index for each word-count pair in the arrays above. xn (ndarray[float]): Numpy array of the total number of words in each document. xm (ndarray[float]): Flattened array of document d_id’s for the lengths above (len = len(x)).

Parameters:

x (Sequence[Tuple[int, Sequence[Tuple[int, float]]]]) – See above for details.

Returns:

See above for details.

Return type:

tuple[Any, tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray]]

row_count(x)[source]

Return the validated number of document-bag rows in an encoded payload.

Parameters:

x (Any)

Return type:

int

multinomial_bag_stream(log_p_vec, min_val, len_dist, combine)[source]

Enumerate normalized integer count-vector bags in descending probability order.

Each fixed-size stream includes log(n!) - sum_w log(c_w!). A real length law is required because a conditional multinomial law for every possible length has infinite total mass when joined without length probabilities.

bag_stream(element_stream, len_dist, combine)[source]

Enumerate bags (multisets) drawn from a sorted element stream, in descending bag-score order.

element_stream is a descending (value, log_prob) iterator over the element distribution (any enumerable element distribution – a fixed categorical, or e.g. a per-document/per-given mixture). A bag scores by the sum of its elements’ log-probs plus log P_len(n) from len_dist; bags enumerate by the per-size multiset best-first search (MultisetProductEnumerator) under a length frontier. When len_dist is Null there is no length term and a synthetic n*log p_max frontier orders the (countably infinite) support by the element term alone. combine maps the tuple of (value, count) pairs to the emitted bag.

class JointMixtureDistribution(components1, components2, w1=None, w2=None, taus12=None, taus21=None, keys=(None, None, None), name=None, joint_weights=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Joint mixture distribution over paired observations.

Observations are (x1, x2) tuples. The first tuple element is scored by components1 and the second by components2.

Parameters:

  • components1 (Sequence[SequenceEncodableProbabilityDistribution])

  • components2 (Sequence[SequenceEncodableProbabilityDistribution])

  • w1 (Sequence[float] | np.ndarray | None)

  • w2 (Sequence[float] | np.ndarray | None)

  • taus12 (list[list[float]] | np.ndarray | None)

  • taus21 (list[list[float]] | np.ndarray | None)

  • keys (tuple[str | None, str | None, str | None] | None)

  • name (str | None)

  • joint_weights (list[list[float]] | np.ndarray | None)

compute_capabilities()[source]

Intersect generated-compute backend support across all child components.

compute_declaration()[source]

Return the generated-compute declaration for the paired latent mixture.

density(x)[source]

Evaluate the density of a joint mixture observation x.

See log_density() for details.

Parameters:

x (Tuple[T0, T1]) – A single (X1, X2) observation.

Returns:

Density evaluated at x.

Return type:

float

log_density(x)[source]

Evaluate the log-density of a joint mixture observation x.

The log-density at x = (x1, x2) is

log(sum_{i=1}^{N} w_i * f_i(x1) * sum_{j=1}^{M} tau12_{ij} * g_j(x2)),

evaluated with a log-sum-exp for numerical stability.

Parameters:

x (Tuple[T0, T1]) – A single (X1, X2) observation.

Returns:

Log-density evaluated at x.

Return type:

float

seq_log_density(x)[source]

Vectorized evaluation of the log-density for an encoded sequence of observations x.

Encoded sequence ‘x’ is a Tuple of length 3 containing:

x[0] (int): Number of observations. x[1] (E0): Encoded sequence of X1 values. x[2] (E1): Encoded sequence of X2 values.

Parameters:

x (tuple[int, E0, E1]) – Encoded sequence of iid joint mixture observations.

Returns:

Log-density evaluated at each observation in the encoded sequence x.

Return type:

ndarray

backend_seq_log_density(x, engine)[source]

Engine-neutral log-density for encoded joint-mixture observations.

Parameters:
Return type:

Any

to_fisher(**kwargs)[source]

Structural Fisher view for the joint mixture.

density_semantics()[source]

Return exact-or-approximate density semantics joined from child components.

sampler(seed=None)[source]

Return a sampler for iid draws from this distribution.

Parameters:

seed (int | None) – Optional random seed.

Returns:

A configured JointMixtureSampler.

Return type:

JointMixtureSampler

estimator(pseudo_count=None)[source]

Return an estimator initialized from this distribution’s components.

Parameters:

pseudo_count (float | None) – Optional smoothing count for latent-state counts.

Returns:

A JointMixtureEstimator.

Return type:

JointMixtureEstimator

dist_to_encoder()[source]

Return an encoder for paired joint-mixture observations.

Return type:

DataSequenceEncoder

enumerator()[source]

Return an enumerator over pairs in descending probability order.

Return type:

JointMixtureEnumerator

class JointMixtureEnumerator(dist)[source]

Bases: DistributionEnumerator

Enumerates the support of a JointMixtureDistribution in descending probability order.

Parameters:

dist (JointMixtureDistribution)

class JointMixtureSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for paired observations from a joint mixture distribution.

Parameters:

  • dist (JointMixtureDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw iid (X1, X2) samples from the joint mixture.

The X1 component state is drawn from w1, X1 is sampled from that component, the X2 component state is drawn from taus12 given the X1 state, and X2 is sampled from the corresponding X2 component.

Parameters:

  • size (int | None) – Number of iid samples to draw. None returns a scalar pair.

  • batched (bool)

Returns:

A scalar pair when size is None; otherwise a list of pairs.

Return type:

tuple[Any, Any] | Sequence[tuple[Any, Any]]

JointMixtureAccumulator

alias of JointMixtureEstimatorAccumulator

JointMixtureAccumulatorFactory

alias of JointMixtureEstimatorAccumulatorFactory

class JointMixtureEstimator(estimators1, estimators2, suff_stat=None, pseudo_count=None, keys=(None, None, None), name=None)[source]

Bases: ParameterEstimator

Estimator for paired latent mixture distributions.

Parameters:

  • estimators1 (Sequence[ParameterEstimator])

  • estimators2 (Sequence[ParameterEstimator])

  • suff_stat (tuple[np.ndarray, np.ndarray, np.ndarray, tuple[E0, ...], tuple[E1, ...]] | None)

  • pseudo_count (float | tuple[float, float, float] | None)

  • keys (tuple[str | None, str | None, str | None] | None)

  • name (str | None)

accumulator_factory()[source]

Return an accumulator factory matching this estimator.

Return type:

JointMixtureEstimatorAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate a joint mixture distribution from aggregated sufficient statistics.

The tuple is interpreted as (x1_counts, x2_counts, joint_counts, x1_child_stats, x2_child_stats).

Parameters:

  • nobs – Weighted number of observations, accepted for the estimator interface.

  • suff_stat (tuple[ndarray, ndarray, ndarray, tuple[E0, ...], tuple[E1, ...]]) – Aggregated joint-mixture sufficient statistics.

Returns:

A fitted joint mixture distribution.

Return type:

JointMixtureDistribution

class JointMixtureDataEncoder(encoder1, encoder2)[source]

Bases: DataSequenceEncoder

Encode paired observations for vectorized joint-mixture scoring and EM.

Parameters:

  • encoder1 (DataSequenceEncoder)

  • encoder2 (DataSequenceEncoder)

seq_encode(x)[source]

Encode a sequence of iid joint mixture observations for vectorized functions.

Return value ‘rv’ is a Tuple containing:

rv[0] (int): Number of observations. rv[1] (E0): Encoded sequence of X1 values. rv[2] (E1): Encoded sequence of X2 values.

Parameters:

x (Sequence[Tuple[T0, T1]]) – Sequence of (X1, X2) observations.

Returns:

See above for details.

Return type:

tuple[int, Any, Any]

class JointMixtureFisherView(dist)[source]

Bases: MixtureFisherView

Complete-data Fisher view for joint mixtures without concrete proxies.

Parameters:

dist (Any)

property num_pairs: int

Number of positive-weight component pairs represented by this Fisher view.

log_density(x)[source]

Evaluate the Fisher-view mixture log-density for one paired observation.

Parameters:

x (Any)

Return type:

float

structured_statistics(x, estimate=None, weight=1.0)[source]

Return posterior pair weights and child sufficient statistics for one observation.

Parameters:
Return type:

Any

class LabeledLDADistribution(topics, alphas, set_dist=None, len_dist=None, gamma_threshold=1.0e-8, max_gamma_iter=1000, fit_diagnostics=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Labeled LDA model for documents with label sets.

Compatible with data type Tuple[Sequence[Tuple[T, float]], Sequence[int]], where T is the data type of the topic distributions.

density(x)[source]

Returns the density (exp of the variational lower bound) for a labeled document x.

Parameters:

x (Tuple[Sequence[Tuple[T, float]], Sequence[int]]) – Document as (value, count) pairs and a list of label indices.

Returns:

Density value for the document x.

density_semantics()[source]

Return density semantics for the labeled-LDA variational bound.

log_density(x)[source]

Returns the variational lower bound (ELBO) on the log-density for a labeled document x.

Parameters:

x (Tuple[Sequence[Tuple[T, float]], Sequence[int]]) – Document as (value, count) pairs and a list of label indices.

Returns:

Lower bound on the log-density of the document x.

seq_log_density(x)[source]

Vectorized evaluation of the variational lower bound (ELBO) for encoded documents.

Arg ‘x’ is the output of ‘LabeledLDADataEncoder.seq_encode()’.

Parameters:

x – Encoded sequence of iid LabeledLDA observations (see LabeledLDADataEncoder.seq_encode()).

Returns:

Numpy array with one lower-bound value per encoded document.

seq_encode(x)[source]

Deprecated alias for dist_to_encoder().seq_encode(): encode iid LabeledLDA observations.

Use dist_to_encoder() and LabeledLDADataEncoder.seq_encode() instead.

Parameters:

x (Sequence[Tuple[Sequence[Tuple[T, float]], Sequence[int]]]) – Sequence of labeled documents.

Returns:

Encoded sequence (see LabeledLDADataEncoder.seq_encode()).

seq_component_log_density(x)[source]

Vectorized per-topic log-density evaluation for encoded documents.

Parameters:

x – Encoded sequence of iid LabeledLDA observations (see LabeledLDADataEncoder.seq_encode()).

Returns:

2-d numpy array (num_documents by nTopics) of per-topic document log-densities.

seq_posterior(x)[source]

Vectorized posterior topic weights for encoded documents.

Parameters:

x – Encoded sequence of iid LabeledLDA observations (see LabeledLDADataEncoder.seq_encode()).

Returns:

2-d numpy array (num_documents by nTopics) of normalized posterior topic weights.

compute_capabilities()[source]

Return backend capability metadata for this concrete LabeledLDA instance.

backend_seq_log_density(x, engine)[source]

Backend-neutral labeled-LDA joint variational lower-bound scoring.

enumerator()[source]

Not supported: LabeledLDA’s log_density is a variational lower bound, not the true marginal.

Each document’s score is an ELBO obtained by running per-document variational inference (seq_posterior), not the exact marginal probability, so enumerating “in descending probability order” is ill-defined – the ELBO ranking need not match the true-probability ranking, and the per-document optimum isn’t a closed-form density over a countable support. Use sampler() and the (approximate) log_density directly.

Return type:

DistributionEnumerator

sampler(seed=None)[source]

Create a sampler for labeled LDA documents.

Note: Requires ‘set_dist’ and ‘len_dist’ to be set.

Parameters:

seed (Optional[int]) – Set seed for random sampling.

Returns:

Sampler bound to this distribution.

Return type:

LabeledLDASampler

estimator(pseudo_count=None)[source]

Create an estimator initialized from this labeled LDA distribution.

Parameters:

pseudo_count (Optional[float]) – Used to re-weight sufficient statistics in estimation.

Returns:

Estimator configured with matching topic estimators.

Return type:

LabeledLDAEstimator

dist_to_encoder()[source]

Return a data encoder for iid labeled LDA observations.

class LabeledLDASampler(dist, seed=None)[source]

Bases: DistributionSampler

Sample labeled documents from a LabeledLDA distribution.

Requires ‘dist.set_dist’ (label sets) and ‘dist.len_dist’ (document lengths) to be set.

sample(size=None, *, batched=True)[source]

Draw iid labeled documents from the LabeledLDA model.

If size is None, a single Tuple[List[T], List[int]] is returned containing the sampled document values and its labels. If size > 0, a list of ‘size’ such Tuples is returned.

Parameters:

  • size (Optional[int]) – Number of iid labeled documents to sample.

  • batched (bool)

Returns:

Tuple[List[T], List[int]] or a List of such Tuples depending on arg size.

class LabeledLDALabelSetStats(stats=None)[source]

Bases: object

Sufficient statistics for the coupled alpha update, grouped by distinct document label set.

Maps each distinct label set S (a sorted tuple of unique label indices) to a pair [n_S, m_S], where n_S is the total weight of the documents carrying label set S and m_S is the weighted sum of the per-document expected log topic weights E[log theta] (a vector with one entry per topic) over those documents.

add(label_set, weight, sum_log_p)[source]

Accumulate document weight and summed expected log topic weights for one label set.

Parameters:

  • label_set (Tuple[int, ...]) – Sorted tuple of label indices.

  • weight (float) – Total document weight to add for the label set.

  • sum_log_p (np.ndarray) – Weighted sum of per-document E[log theta] vectors to add.

Returns:

None.

combine(other)[source]

Merge the statistics of another LabeledLDALabelSetStats instance into this instance.

Parameters:

other (LabeledLDALabelSetStats) – Statistics to merge in (left unmodified).

Returns:

LabeledLDALabelSetStats object (self).

copy()[source]

Returns a deep copy of the LabeledLDALabelSetStats instance.

arrays()[source]

Returns the statistics as parallel arrays in sorted label-set order.

Returns:

Tuple of the sorted label-set tuples (List[Tuple[int, …]]), the per-set weights n_S (1-d numpy array), and the per-set summed expected log topic weights m_S (2-d numpy array with one row per label set).

LabeledLDAAccumulator

alias of LabeledLDAEstimatorAccumulator

LabeledLDAAccumulatorFactory

alias of LabeledLDAEstimatorAccumulatorFactory

class LabeledLDAEstimator(estimators, num_alphas, suff_stat=None, pseudo_count=None, keys=(None, None), fixed_alpha=None, gamma_threshold=1.0e-8, alpha_threshold=1.0e-8, max_gamma_iter=1000, max_alpha_iter=2000, set_dist=None, len_dist=None)[source]

Bases: ParameterEstimator

Estimate labeled LDA distributions from aggregated sufficient statistics.

accumulator_factory()[source]

Return an accumulator factory configured from this estimator.

accumulatorFactory()[source]

Deprecated alias for accumulator_factory().

estimate(nobs, suff_stat)[source]

Estimate a labeled LDA distribution from aggregated sufficient statistics.

suff_stat is a tuple containing:

suff_stat[0] (Optional[np.ndarray]): Previous alphas matrix. suff_stat[1] (LabeledLDALabelSetStats): Per-label-set expected log topic weights and counts. suff_stat[2] (Union[float, np.ndarray]): Per-label weighted document counts. suff_stat[3] (np.ndarray): Label-allocated weighted topic counts. suff_stat[4] (Sequence): Sufficient statistics for the topic distribution accumulators.

If ‘fixed_alpha’ is None, the alphas matrix is re-estimated by maximizing the coupled objective over all label rows (see ‘update_alpha_coupled()’). When every document carries exactly one label the objective decouples and the per-row fixed-point updates are used (see ‘update_alpha()’). Otherwise the alphas matrix is set to ‘fixed_alpha’.

Parameters:

  • nobs (Optional[float]) – Number of observations used in estimation.

  • suff_stat – See above for details.

Returns:

Estimated distribution.

Return type:

LabeledLDADistribution

class LabeledLDADataEncoder(encoder, num_alphas=None)[source]

Bases: DataSequenceEncoder

Encode iid labeled LDA observations for vectorized scoring.

seq_encode(x)[source]

Encode labeled documents under the canonical nonempty-label schema.

row_count(x)[source]

Return the validated number of document rows in an encoded payload.

doc_label_sets(nbx, nbcnt)[source]

Returns the sorted label-set tuple of each encoded document.

Parameters:

  • nbx (np.ndarray) – Flattened array of label indices over all documents (document-contiguous).

  • nbcnt (np.ndarray) – Number of labels for each document.

Returns:

List of sorted label-index tuples, one per document.

coupled_alpha_objective(alpha, label_sets, set_counts, set_mean_logs)[source]

Coupled multi-label alpha objective (terms independent of alpha dropped).

F(alpha) = sum_S n_S * [ log Gamma(sum_k a_Sk) - sum_k log Gamma(a_Sk) + sum_k a_Sk * mbar_Sk ], where a_S = mean_{l in S} alpha[l], n_S = set_counts[S], and mbar_S = set_mean_logs[S] are the per-set mean expected log topic weights.

Parameters:

  • alpha (np.ndarray) – Alphas matrix (num_alphas by num_topics).

  • label_sets (Sequence[Tuple[int, ...]]) – Label-set tuples.

  • set_counts (np.ndarray) – Per-set document weights n_S.

  • set_mean_logs (np.ndarray) – Per-set mean expected log topic weights mbar_S (one row per set).

Returns:

Objective value F(alpha).

coupled_alpha_gradient(alpha, label_sets, set_counts, set_mean_logs)[source]

Gradient of the coupled multi-label alpha objective with respect to alpha.

dF/d alpha[l,k] = sum_{S contains l} (n_S/|S|) * [ psi(sum_j a_Sj) - psi(a_Sk) + mbar_Sk ], with one term per occurrence of l in S.

Parameters:

  • alpha (np.ndarray) – Alphas matrix (num_alphas by num_topics).

  • label_sets (Sequence[Tuple[int, ...]]) – Label-set tuples.

  • set_counts (np.ndarray) – Per-set document weights n_S.

  • set_mean_logs (np.ndarray) – Per-set mean expected log topic weights mbar_S (one row per set).

Returns:

Numpy 2-d array with the same shape as alpha.

update_alpha_coupled(current_alpha, label_sets, set_counts, set_mean_logs, alpha_threshold, max_its=2000, *, return_diagnostics=False)[source]

Maximize the validated coupled label-row objective with bounded ascent.

class LDAOptimizationDiagnostics(algorithm, converged, iterations, max_iterations, termination_reason, final_residual, objective_trace=(), residual_trace=(), impossible_documents=())[source]

Bases: object

Serializable termination record for an LDA inner optimization.

Parameters:
exception LDAConvergenceError(diagnostics)[source]

Bases: RuntimeError

Raised when an LDA inner optimization cannot meet its declared contract.

Parameters:

diagnostics (LDAOptimizationDiagnostics)

class LDADistribution(topics, alpha, len_dist=NullDistribution(), gamma_threshold=1.0e-8, max_gamma_iter=100, fit_diagnostics=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Latent Dirichlet allocation model for documents given as bags of weighted values.

Data type: Sequence[Tuple[T, float]], where T is the data type of the topic distributions and each (value, count) pair gives the count of a value in the document.

Parameters:

  • topics (Sequence[SequenceEncodableProbabilityDistribution])

  • alpha (Sequence[float] | ndarray)

  • len_dist (SequenceEncodableProbabilityDistribution | None)

  • gamma_threshold (float)

  • max_gamma_iter (int)

  • fit_diagnostics (LDAOptimizationDiagnostics | None)

compute_capabilities()[source]

Return backend capability metadata for this concrete LDA instance.

compute_declaration()[source]

Return the generated-compute declaration for latent Dirichlet allocation.

density(x)[source]

Evaluate the density of a single LDA document.

See log_density() for details.

Parameters:

x (Sequence[Tuple[int, float]]) – A document given as (value, count) pairs.

Returns:

Density evaluated at x.

Return type:

float

density_semantics()[source]

Return density semantics for the variational LDA document bound.

log_density(x)[source]

Evaluate the log-density of a single LDA document.

Note: The returned value is the variational lower bound (ELBO) on the marginal document log-likelihood obtained from the standard LDA mean-field approximation, not the exact (intractable) marginal log-likelihood.

Parameters:

x (Sequence[Tuple[int, float]]) – A document given as (value, count) pairs.

Returns:

Variational lower bound on the log-density evaluated at x.

Return type:

float

seq_log_density(x)[source]

Vectorized evaluation of the document log-densities for an encoded corpus x.

Encoded sequence ‘x’ is a Tuple of length 5 containing:

x[0] (int): Number of documents in corpus. x[1] (np.ndarray): Document id for flattened array of values. x[2] (np.ndarray): Flattened array of counts for each value in each document. x[3] (Optional[np.ndarray]): Optional warm-start gammas (defaults to None). x[4] (E0): Sequence encoded flattened values.

Note: Returns the per-document variational lower bound (ELBO); see log_density(). If a document-length distribution ‘len_dist’ is set, its log-density of the total token count of each document is added to the returned values.

Parameters:

x (tuple[int, ndarray, ndarray, ndarray | None, E0]) – Encoded corpus of LDA documents (see LDADataEncoder.seq_encode()).

Returns:

Numpy array of log-density (ELBO) values, one entry per document.

Return type:

ndarray

backend_seq_log_density(x, engine)[source]

Backend-neutral LDA variational lower-bound scoring.

Parameters:
Return type:

Any

seq_component_log_density(x)[source]

Vectorized evaluation of the per-topic log-density of each document in encoded corpus x.

Parameters:

x (tuple[int, ndarray, ndarray, ndarray | None, E0]) – Encoded corpus of LDA documents (see LDADataEncoder.seq_encode()).

Returns:

2-d numpy array with shape (number of documents, n_topics), where entry (i, l) is the log-density of document i evaluated entirely under topic l.

Return type:

ndarray

backend_seq_component_log_density(x, engine)[source]

Backend-neutral per-topic document scores.

Parameters:
Return type:

Any

seq_posterior(x)[source]

Vectorized evaluation of the posterior topic proportions for each document in encoded corpus x.

The variational gammas are computed for each document and normalized to sum to one.

Parameters:

x (tuple[int, ndarray, ndarray, ndarray | None, E0]) – Encoded corpus of LDA documents (see LDADataEncoder.seq_encode()).

Returns:

2-d numpy array with shape (number of documents, n_topics) containing posterior topic proportions for each document.

Return type:

ndarray

latent_posterior(doc)[source]

Return the mean-field variational posterior q(theta, z) for a single document.

Runs the per-document Blei-Ng-Jordan variational fixed point and returns a MeanFieldLDAPosterior: .topic_proportions() (the document-topic mix E[theta]), .marginals() (per-word topic responsibilities phi), .sample(rng) (theta, z), .mode() (MAP topic per word), or .entropy().

Parameters:

doc (Sequence[tuple[int, float]])

Return type:

MeanFieldLDAPosterior

posterior_predictive(doc, n_words, seed=None)[source]

Draw n_words new words conditioned on the document doc.

Sample the document-topic mix theta ~ q(theta) = Dir(gamma) from the variational posterior, then generate each new word by drawing a topic ~ theta and a word from that topic – “given this document, generate more words from its inferred topic mixture”.

Parameters:
Return type:

list[Any]

sampler(seed=None)[source]

Create a sampler for documents from this distribution.

Parameters:

seed (Optional[int]) – Seed for the random number generator used in sampling.

Returns:

Sampler bound to this distribution.

Return type:

LDASampler

estimator(pseudo_count=None)[source]

Create an estimator initialized from this distribution’s topics.

Parameters:

pseudo_count (Optional[float]) – If passed, used to re-weight sufficient statistics during estimation.

Returns:

Estimator configured with matching topic and length estimators.

Return type:

LDAEstimator

dist_to_encoder()[source]

Return a data encoder for iid LDA documents.

Return type:

LDADataEncoder

enumerator()[source]

LDA does not support enumeration.

The document log-density is a variational lower bound (ELBO) over latent topic assignments rather than an exact density, so an enumeration satisfying log_prob == log_density over a well-defined support cannot be constructed.

Raises:

EnumerationError – Always.

Return type:

DistributionEnumerator

class LDASampler(dist, seed=None)[source]

Bases: DistributionSampler

Sample documents from an LDA distribution.

Parameters:

  • dist (LDADistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw one or ‘size’ documents from the LDA model.

Note: Sample return value is not counted by value! Each document is returned as a flat list of sampled topic values (use mixle.utils.optsutil.count_by_value to obtain (value, count) pairs).

With batched=True (default), when size is not None the per-document lengths, Dirichlet proportions and topic-count multinomials are drawn first, then every token across the whole batch is grouped by topic and each topic sampler is invoked once. Because the topic samplers are consumed in topic order rather than per-document order, the token draws are statistically equivalent but NOT byte-identical to batched=False. The length, Dirichlet and multinomial draws are byte-identical (same order). Set batched=False to reproduce the exact legacy per-document output for a given seed.

Parameters:

  • size (Optional[int]) – Number of documents to sample. If None, a single document is returned.

  • batched (bool) – Vectorize token draws across documents (default); set False for the legacy per-document loop.

Returns:

A single document (list of values) if size is None, else a list of ‘size’ documents.

Return type:

Sequence[Any] | Any

LDAAccumulator

alias of LDAEstimatorAccumulator

LDAAccumulatorFactory

alias of LDAEstimatorAccumulatorFactory

class LDAEstimator(estimators, len_estimator=NullEstimator(), suff_stat=None, pseudo_count=None, keys=(None, None), fixed_alpha=None, gamma_threshold=1.0e-8, alpha_threshold=1.0e-8, max_gamma_iter=100, max_alpha_iter=1000)[source]

Bases: ParameterEstimator

Estimate LDA distributions from aggregated variational sufficient statistics.

Parameters:
accumulator_factory()[source]

Return an accumulator factory configured from this estimator.

Return type:

LDAEstimatorAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate an LDA distribution from aggregated sufficient statistics.

suff_stat is a six-item tuple containing:

suff_stat[0] (Optional[np.ndarray]): Previous Dirichlet parameter estimate. suff_stat[1] (np.ndarray): Aggregated expected log topic proportions. suff_stat[2] (float): Aggregated weighted document count. suff_stat[3] (np.ndarray): Aggregated weighted per-topic value counts. suff_stat[4] (Sequence[SS0]): Sufficient statistics for each topic. suff_stat[5] (Optional[Any]): Sufficient statistics for the document-length distribution.

Parameters:

  • nobs (Optional[float]) – Weighted number of observations used in aggregation of suff_stat.

  • suff_stat – See above for details.

Returns:

Estimated distribution.

Return type:

LDADistribution

class LDADataEncoder(encoder)[source]

Bases: DataSequenceEncoder

Encode iid LDA documents for vectorized scoring.

Parameters:

encoder (DataSequenceEncoder)

seq_encode(x)[source]

Encode a sequence of iid LDA observations for vectorized functions.

Return value ‘rv’ is a Tuple containing:

rv[0] (int): Number of documents in corpus. rv[1] (np.ndarray): Document id for flattened array of values. rv[2] (np.ndarray): Flattened array of counts for each value in each document. rv[3] (Optional[np.ndarray]): Currently default to None rv[4] (E0): Sequence encoded flattened values.

Parameters:

x (Sequence[Sequence[Tuple[int, float]]]) – Sequence of LDA documents.

Returns:

See above for details.

Return type:

tuple[int, ndarray, ndarray, Any | None, Any]

row_count(x)[source]

Return the validated number of document rows in an encoded corpus.

Parameters:

x (Any)

Return type:

int

class LookbackHiddenMarkovModelDistribution(topics, w=MISSING, transitions=MISSING, lag=0, init_dist=None, len_dist=NullDistribution(), name=None, weights=MISSING, terminal_states=None, prior=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Hidden Markov model whose state emissions condition on the previous lag observations.

Parameters:

  • topics (Sequence[SequenceEncodableProbabilityDistribution])

  • w (ndarray)

  • lag (int)

  • init_dist (Sequence[SequenceEncodableProbabilityDistribution] | None)

  • len_dist (SequenceEncodableProbabilityDistribution | None)

  • name (str | None)

  • weights (ndarray)

  • terminal_states (set[int] | Sequence[int] | None)

get_prior()[source]

Returns the chain conjugate prior in (init_prior, row_priors) form (or None).

Per-state emission/init component priors are owned by those distributions themselves – mirrors get_prior().

set_prior(prior)[source]

Set the conjugate Dirichlet chain prior over w/transitions and flag has_conj_prior.

Mirrors set_prior() exactly (no digamma-expectation caching: this class has no expected_log_density counterpart to feed). prior=None leaves the distribution a plain point model.

Parameters:

prior – (init_prior, row_priors) tuple or None.

Return type:

None

density(x)[source]

Evaluate the density of the distribution at sequence x.

Parameters:

x (Sequence[T]) – Observed sequence.

Returns:

Density at x.

Return type:

float

log_density(x)[source]

Evaluate the log-density of the distribution at sequence x.

Marginalizes the hidden state path with a scaled forward pass. The initial segment x[:lag] is scored by init_dist, each window x[t-lag:t+1] by the topic distributions, and the number of hidden positions by len_dist. When lag == 0 there is no initial segment: the first state emits the window x[0:1] directly (ordinary HMM).

Parameters:

x (Sequence[T]) – Observed sequence with len(x) >= lag.

Returns:

Log-density at x.

Return type:

float

viterbi_sequence(x)[source]

Compute the most likely hidden state sequence for observed sequence x.

Parameters:

x (Sequence[T]) – Observed sequence with len(x) >= lag.

Returns:

Integer array of len(x) - lag + 1 (len(x) when lag == 0) most likely hidden

state indices.

Return type:

np.ndarray

seq_log_density(x)[source]

Vectorized evaluation of the log-density at encoded sequences x.

Parameters:

x – Encoded sequence data produced by seq_encode() / dist_to_encoder().

Returns:

Log-density value for each encoded sequence.

Return type:

np.ndarray

compute_capabilities()[source]

Return backend capability metadata for this concrete lookback-HMM instance.

backend_seq_log_density(x, engine)[source]

Engine-neutral lookback-HMM scoring via the shared HMM forward pass.

seq_posterior(x, filtered=False)[source]

Return per-sequence posterior state probabilities for encoded observations.

Each returned array holds the forward-backward SMOOTHING marginals P(state_t = i | the whole sequence), matching every other seq_posterior in the package – including HiddenMarkovModelDistribution.seq_posterior(), whose docstring already stated that contract on behalf of both. Pass filtered=True for the forward-only P(state_t = i | observations up to t).

This ran the forward-only kernel unconditionally, so it returned the FILTERED probabilities under the smoothed name: a different quantity that agrees at the last position and nowhere else (on a four-step categorical chain, off by 0.16 at position 0). Both are legitimate to want and only one is a posterior over the observed sequence, which is what the method is named for (R05-F04).

Parameters:

  • x – Encoded sequence data produced by seq_encode() / dist_to_encoder().

  • filtered (bool) – return the forward-only filtered probabilities instead of the smoothed ones.

Returns:

For each sequence, an array of per-position posterior state

probabilities with shape (num_windows, num_states).

Return type:

List[np.ndarray]

Raises:

NotImplementedError – when the model restricts terminal_states. The recursion below is the unrestricted one, and running it under that restriction would answer for a different model – see the refusal in the body.

density_semantics()[source]

Return exact-or-approximate density semantics joined from child models.

sampler(seed=None)[source]

Create a LookbackHiddenMarkovModelSampler for this distribution.

Parameters:

seed (Optional[int]) – Seed for random number generator.

Returns:

Sampler object (requires a non-null len_dist).

Return type:

LookbackHiddenMarkovModelSampler

enumerator()[source]

Not supported: lookback emissions condition on the previous lag observations.

Unlike the standard / segmental HMM, each emission depends on the preceding lag symbols, so the effective state is (hidden_state, last lag observations). Best-first enumeration over that augmented (and, for large/continuous alphabets, unbounded) state space is not implemented; use sampler() or the exact log_density / viterbi_sequence instead.

Return type:

DistributionEnumerator

estimator(pseudo_count=None)[source]

Create a LookbackHiddenMarkovModelEstimator from this distribution.

Parameters:

pseudo_count (Optional[float]) – Regularize the initial-state and transition estimates.

Returns:

Estimator built from the topic, initial-segment, and

length distributions, preserving the lag.

Return type:

LookbackHiddenMarkovModelEstimator

seq_encode(x)[source]

Encode a sequence of observed sequences for vectorized seq_ calls.

Parameters:

x (Sequence[Sequence[T]]) – Sequence of iid observed sequences.

Returns:

Encoded data consistent with seq_log_density(), seq_posterior(), and seq_update().

dist_to_encoder()[source]

Return a LookbackHiddenMarkovModelDataEncoder for encoding sequences of iid observations.

Returns:

Encoder built from the topic, initial, and length

distributions of this instance.

Return type:

LookbackHiddenMarkovModelDataEncoder

class LookbackHiddenMarkovModelSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for LookbackHiddenMarkovModelDistribution. Requires non-null init_dist and len_dist.

Parameters:

  • dist (LookbackHiddenMarkovModelDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw iid sequences from the lookback hidden Markov distribution.

Parameters:

  • size (Optional[int]) – Number of sequences to draw. If None, a single sequence is returned.

  • batched (bool)

Returns:

One sampled sequence if size is None, else a list of

size sampled sequences.

Return type:

Union[List[T], List[List[T]]]

LookbackHiddenMarkovModelAccumulator

alias of LookbackHiddenMarkovModelEstimatorAccumulator

LookbackHiddenMarkovModelAccumulatorFactory

alias of LookbackHiddenMarkovModelEstimatorAccumulatorFactory

class LookbackHiddenMarkovModelEstimator(estimators, lag=0, init_estimators=None, len_estimator=NullEstimator(), suff_stat=None, pseudo_count=(None, None), name=None, keys=(None, None, None), terminal_states=None, prior=None)[source]

Bases: ParameterEstimator

Estimator for a lookback hidden Markov model from aggregated sufficient statistics.

Parameters:

  • estimators (Sequence[ParameterEstimator])

  • lag (int)

  • init_estimators (Sequence[ParameterEstimator] | None)

  • len_estimator (ParameterEstimator | None)

  • pseudo_count (float | tuple[float | None, float | None] | None)

  • name (str | None)

  • keys (tuple[str | None, str | None, str | None] | None)

get_prior()[source]

Returns the chain conjugate prior in (init_prior, row_priors) form (or None).

Mirrors get_prior().

set_prior(prior)[source]

Set the conjugate Dirichlet chain prior and flag whether it admits the conjugate update.

Mirrors set_prior() exactly.

Parameters:

prior – (init_prior, row_priors) tuple or None; has_conj_prior is set when both the initial-state prior and all row priors are Dirichlet.

Return type:

None

model_log_density(model)[source]

Log-density of the model parameters under the priors (ELBO global term).

Sums the Dirichlet log-densities of the initial-state and transition probabilities (when a conjugate chain prior is set) plus each topic/init-segment estimator’s own model_log_density of its fitted distribution. Mirrors model_log_density() (which does not include its own len_estimator either – kept at parity here rather than widening scope beyond the sibling class).

Parameters:

model (LookbackHiddenMarkovModelDistribution) – Model to score.

Returns:

Prior log-density of the model parameters.

Return type:

float

accumulator_factory()[source]

Create a LookbackHiddenMarkovModelEstimatorAccumulatorFactory from the member estimators.

Returns:

Factory for accumulators consistent

with this estimator.

Return type:

LookbackHiddenMarkovModelEstimatorAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate a LookbackHiddenMarkovModelDistribution from aggregated sufficient statistics.

Parameters:

  • nobs (Optional[float]) – Weighted number of observations (passed to the length estimator).

  • suff_stat (Tuple) – Sufficient statistics in the format returned by LookbackHiddenMarkovModelEstimatorAccumulator.value().

Returns:

M-step estimate of the distribution.

Return type:

LookbackHiddenMarkovModelDistribution

class LookbackHiddenMarkovModelDataEncoder(encoder, lag, len_encoder=NullDataEncoder(), init_encoder=NullDataEncoder())[source]

Bases: DataSequenceEncoder

Encoder for sequences of iid lookback-HMM observations (each a Sequence[T]).

Parameters:

  • encoder (DataSequenceEncoder)

  • lag (int)

  • len_encoder (DataSequenceEncoder | None)

  • init_encoder (DataSequenceEncoder | None)

seq_encode(x)[source]

Encode a sequence of iid observed sequences for vectorized processing.

Each sequence x[i] is split into its initial segment x[i][:lag] and its sliding windows x[i][j-lag:j+1] for j in [lag, len(x[i])); index arrays track which flattened position belongs to which sequence and which positions are initial segments vs emission windows. When lag == 0 there are no initial segments: every position is an emission window (ordinary HMM).

Parameters:

x (Sequence[Sequence[T]]) – Sequence of iid observed sequences.

Returns:

((ids, idi, ims, imi, sz, enc_windows, enc_inits), len_enc) where ids/idi map

windows/initial segments to sequence indices, ims/imi give their flattened positions, sz holds per-sequence position counts, enc_windows/enc_inits are the encoded windows and initial segments (enc_inits is None when lag == 0), and len_enc is the encoded position counts.

Return type:

Tuple

row_count(x)[source]

Return and validate the number of encoded lookback-HMM sequences.

Parameters:

x (Any)

Return type:

int

class LookbackTerminalDataEncoder[source]

Bases: DataSequenceEncoder

Passthrough encoder for terminal-state lookback HMMs: keeps raw sequences (scored per sequence).

seq_encode(x)[source]

Encode terminal-aware lookback HMM sequences as lists.

exception HiddenMarkovNonterminationError(*, mode, max_steps, last_state=None)[source]

Bases: RuntimeError

Raised when bounded stopping-time sampling does not reach its terminal event.

Parameters:

  • mode (str)

  • max_steps (int)

  • last_state (int | None)

Return type:

None

validated_state_ids(values, n_states, *, context, field)[source]

Return exact in-range state IDs without lossy integer coercion.

Parameters:
Return type:

frozenset[int] | None

validated_terminal_states(values, n_states, *, context)[source]

Return exact in-range terminal state IDs without lossy integer coercion.

Parameters:
Return type:

frozenset[int] | None

validated_terminal_values(values, *, context)[source]

Return an owned non-empty set of terminal emissions.

Parameters:
Return type:

frozenset[Any] | None

validate_terminal_reachability(initial, transitions, terminal_states, *, context)[source]

Require every initially reachable state to have a positive-probability path to termination.

Parameters:
Return type:

None

validated_terminal_step_cap(value)[source]

Return an exact positive stopping-time sampling bound.

Parameters:

value (Any)

Return type:

int

require_terminal_reached(reached, *, mode, max_steps, last_state=None)[source]

Raise the typed bounded-sampling failure when the stop event was not reached.

Parameters:

  • reached (bool)

  • mode (str)

  • max_steps (int)

  • last_state (int | None)

Return type:

None

mixture_prior(weight_prior, component_priors)[source]

Build the joint mixture prior: a weight prior plus one prior per component.

Parameters:

  • weight_prior (SequenceEncodableProbabilityDistribution) – Prior on the mixture weights (a DirichletDistribution or SymmetricDirichletDistribution).

  • component_priors (Sequence[SequenceEncodableProbabilityDistribution]) – Sequence of one conjugate prior per component.

Returns:

A (weight_prior, tuple(component_priors)) pair consumed by MixtureDistribution/MixtureEstimator set_prior.

Return type:

tuple[SequenceEncodableProbabilityDistribution, tuple[SequenceEncodableProbabilityDistribution, …]]

class MixtureDistribution(components, w=MISSING, name=None, weights=MISSING, prior=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Finite mixture over homogeneous component distributions.

components define both the conditional families p(x | z=k) and the observation type accepted by the mixture. w contains the component probabilities and is cached as log_w for stable scalar and vectorized scoring. Zero-weight components are retained for structural compatibility but contribute -inf to log-density calculations.

Parameters:

  • components (Sequence[SequenceEncodableProbabilityDistribution]) – Component distributions. Each component should support the same raw observation shape and sequence-encoding contract.

  • w (np.ndarray | list[float]) – Component weights. The values are interpreted as simplex weights and should sum to one.

  • name (str | None) – Optional display name for diagnostics and generated artifacts.

  • weights (np.ndarray | list[float]) – Alias for w.

  • prior (SequenceEncodableProbabilityDistribution | None) – Optional joint mixture prior or weight prior.

components

Component distribution objects.

w

Component weights as a NumPy array.

zw

Boolean mask for zero-weight components.

log_w

Log weights, with zero-weight entries represented as -inf.

num_components

Number of mixture components.

compute_capabilities()[source]

Return compute-backend metadata shared by all mixture components.

compute_declaration()[source]

Return the symbolic declaration for mixture weights and component statistics.

get_prior()[source]

Return the joint mixture prior, or None for a plain point model.

When a weight prior is attached the joint prior is the (weight_prior, tuple(component priors)) pair produced by mixture_prior(); otherwise None.

Return type:

SequenceEncodableProbabilityDistribution | None

set_prior(prior)[source]

Attach a weight prior (and optional per-component priors), caching weight expectations.

With a (symmetric) Dirichlet weight prior this caches the variational weight expectations E[log w_k] = digamma(alpha_k) - digamma(sum_j alpha_j) used by expected_log_density. Component priors, when supplied, are delegated to each component via component.set_prior. prior=None (the default) leaves the mixture a plain point model (byte-identical MLE behaviour).

Parameters:

prior (SequenceEncodableProbabilityDistribution | None)

Return type:

None

expected_log_density(x)[source]

Return the optimized local variational log-evidence bound.

This compatibility name does not compute the generally intractable E[log(sum_k w_k p_k(x))]. With conjugate priors it returns logsumexp_k(E[log w_k] + E[log p_k(x)]), the optimized local variational lower bound. Use variational_log_evidence_bound() when the semantic distinction matters. Without a conjugate weight prior it falls back to the plug-in log_density.

Parameters:

x (T)

Return type:

float

variational_log_evidence_bound(x)[source]

Return the local variational lower bound for one observation.

Parameters:

x (T)

Return type:

float

seq_expected_log_density(x)[source]

Return vectorized local variational log-evidence bounds.

This is the sequence form of variational_log_evidence_bound(), retained under the framework-wide compatibility name.

Parameters:

x (T1)

Return type:

ndarray

seq_variational_log_evidence_bound(x)[source]

Return local variational lower bounds for encoded observations.

Parameters:

x (T1)

Return type:

ndarray

density(x)[source]

Return the mixture density at one raw observation.

Parameters:

x (T) – Observation accepted by every component family.

Returns:

exp(log_density(x)).

Return type:

float

density_semantics()[source]

Return joined density semantics over all mixture components.

log_density(x)[source]

Return the mixture log-density at one raw observation.

The calculation is logsumexp_k(log p_k(x) + log w_k). Component families are responsible for returning -inf for observations outside support; the mixture combines those values without converting them to NaN.

Parameters:

x (T) – Observation accepted by every component family.

Returns:

Finite log-density when at least one positive-weight component can score the observation, otherwise -inf.

Return type:

float

conditional(observed)[source]

Return the conditional mixture over the unobserved coordinates given observed.

The conditional of a mixture is itself a mixture: for sum_k w_k f_k observing x_o,

P(x_u | x_o) = sum_k w’_k f_k(x_u | x_o), w’_k proportional to w_k f_k.marginal(x_o)(x_o),

i.e. the component responsibilities are updated by how well each component explains the observed coordinates and each component is replaced by its own conditional. Because the result is a full MixtureDistribution you can both score it and .sampler(seed).sample() from it – the latter is given=-style conditional sampling that first draws a component from the posterior responsibilities, then draws the unobserved coordinates from that component’s conditional.

Requires each component to support marginal(indices) and condition(observed) (e.g. the multivariate Gaussian / Student-t). observed maps coordinate index to its fixed value.

Parameters:

observed (dict[int, float])

Return type:

MixtureDistribution

component_log_density(x)[source]

Return component-wise log densities for one raw observation.

Parameters:

x (T) – Observation accepted by every component family.

Returns:

One log-density per component, before mixture weights are applied.

Return type:

ndarray

posterior(x)[source]

Return component responsibilities for one raw observation.

Responsibilities are proportional to w[k] * p_k(x). An observation outside every component’s support receives zero responsibility and therefore cannot fabricate component statistics.

Parameters:

x (T) – Observation accepted by every component family.

Returns:

Responsibility vector over component labels. It sums to one when the evidence is possible and is all zero otherwise.

Return type:

ndarray

seq_component_log_density(x)[source]

Return vectorized component log densities for encoded observations.

x must be produced by MixtureDataEncoder.seq_encode or by an equivalent component encoder. The output has shape (n, k) where n is the number of encoded observations and k is the number of mixture components.

Parameters:

x (T1) – Encoded observation batch.

Returns:

Component log-density matrix before mixture weights are applied.

Return type:

ndarray

seq_log_density(x)[source]

Return vectorized mixture log densities for encoded observations.

Each row is evaluated with a row-wise log-sum-exp over component scores plus log weights. Rows for which every positive-weight component is impossible return -inf.

Parameters:

x (T1) – Encoded observation batch.

Returns:

One log-density per encoded observation.

Return type:

ndarray

backend_seq_component_log_density(x, engine)[source]

Engine-neutral component log densities for encoded data.

Parameters:

  • x (T1)

  • engine (Any)

Return type:

Any

backend_seq_log_density(x, engine)[source]

Engine-neutral mixture log-density for encoded data.

Parameters:

  • x (T1)

  • engine (Any)

Return type:

Any

gradient_fit_state(engine, torch, leaves, recurse, tensor_param)[source]

Return distribution-owned state for autograd fitting.

Parameters:
Return type:

Any

seq_posterior(x)[source]

Return vectorized component responsibilities for encoded observations.

Each row is proportional to w[k] * p_k(x_i). Rows where all positive-weight components are impossible receive zero responsibility, matching posterior() without inventing latent assignments.

Parameters:

x (T1) – Encoded observation batch.

Returns:

(n, k) responsibility matrix. Possible rows sum to one; impossible rows sum to zero.

Return type:

ndarray

latent_posterior(x)[source]

Return the latent posterior q(z | x) over component labels for raw observations x.

q(z) is the exact independent-categorical posterior whose marginals are the EM responsibilities. The returned CategoricalLatentPosterior can .marginals() (the responsibilities), .sample(rng) component labels, .mode() (the MAP labels), or .entropy().

Parameters:

x (Sequence[T])

Return type:

CategoricalLatentPosterior

posterior_predictive(x, seed=None)[source]

Draw posterior-predictive observations conditioned on x.

For each observed x_i the component is sampled from the latent posterior q(z_i | x_i) and a fresh observation is emitted from that component – i.e. “given I saw x_i, draw a new point from the same mixture component it likely came from”. Returns a list the length of x. Draws are grouped by component and scattered (vectorized) via the shared sampling helper.

Parameters:
Return type:

list[Any]

support_size()[source]

Upper bound on distinct support points: the sum over components (union <= sum).

Return type:

int | None

tropical_displacement_bits()[source]

log2(#positive-weight components) – the tropical-vs-marginal cost gap (in bits).

The marginal log p(x) = logsumexp_k (log w_k + log p_k(x)) is bounded by its largest term M(x) = max_k (log w_k + log p_k(x)) via M(x) <= log p(x) <= M(x) + log K, where K is the number of components that can contribute (positive weight). The structural seek bins by the tropical cost M(x); mixle.enumeration.density_rank.marginal_seek() widens its smear window by this many bits so the reported rank bracket provably contains the TRUE marginal rank. K <= 1 means the marginal is a single term -> 0.0 (the seek is then exact). When the component supports are provably disjoint every value lands in one component, so M(x) equals the marginal and there is likewise no displacement -> 0.0 (the seek is exact and tight).

Return type:

float

to_fisher(**kwargs)[source]

Structural Fisher view for the mixture.

sampler(seed=None)[source]

Return a sampler that draws from the mixture distribution.

Parameters:

seed (int | None) – Optional RandomState seed for reproducible draws.

Returns:

MixtureSampler bound to this distribution.

Raises:

ValueError – If this mixture is a likelihood factor rather than a normalized law.

Return type:

MixtureSampler

A component may be a scaled law – an unnormalized categorical, an open-world smoothed emission, a component that won zero responsibility – and _owned_generative_components admits those deliberately, because a constant scale cancels in the E-step’s responsibility normalization. It does NOT cancel in the composite: a mass-0.75 component under weight 0.5 leaves the mixture with total mass 0.875, and density_semantics() reports LIKELIHOOD_FACTOR to say so. The sampler used to ignore that and draw as though the weights described a law, so the scorer and the sampler described different objects (MXR-080-1857).

Refusing here rather than at construction is the point: scoring such a mixture is legitimate and tested (an open-world component has infinite total mass and no finite scale to absorb), so the object must remain constructible. What it cannot do is pretend to be a generative law.

estimator(pseudo_count=None)[source]

Return an estimator with matching component structure.

Parameters:

pseudo_count (float | None) – Optional smoothing mass applied through the estimator path.

Returns:

MixtureEstimator suitable for fitting observations of the same type as this distribution.

Return type:

MixtureEstimator

decomposition()[source]

Mixture components split along the component axis. Responsibilities (logsumexp) are computed INSIDE a shard; across shards the per-component sufficient stats SUM-reduce plus one scalar total-count all-reduce – the homogeneous stacked-kernel + DTensor path (engine_axis=0).

dist_to_encoder()[source]

Return an encoder that delegates observation encoding to components.

Return type:

MixtureDataEncoder

enumerator()[source]

Return an enumerator over the union of component supports.

Return type:

MixtureEnumerator

quantized_index(max_bits, bin_width_bits=1.0)[source]

Build a bounded bit-quantized index from a global mixture frontier.

The primary path pulls candidates from weighted component enumerator heads. The log-sum of those heads bounds every unseen value, so construction stops when the live global frontier falls below 2**(-max_bits). This avoids the looser per-component log2(K) candidate expansion. If a component cannot enumerate, the method falls back to the structured cross-index path.

Parameters:
Return type:

QuantizedEnumerationIndex

quantized_count_index(quantizer, max_fine_bucket)[source]

BoundedCount for the MARGINAL mixture law: pool weight-scaled component count indices.

log p(x) = logsumexp_k (log w_k + log p_k(x)) has no exact structural count – overlapping component supports would need value-level deduplication. This builds the count semiring’s plus-fold over scale(component_index, log w_k) instead, which:

  • reaches a 2**M budget structurally (no enumeration), and

  • is a conservative UPPER bound – a value shared by several components is counted once per component, and each value is binned by its dominant weighted component (the tropical cost, within log2(K) bits of the exact logsumexp).

Every unranked value still carries its exact mixture log_density (re-evaluated by the budget builder). For an exact small-budget index (best-first union with dedup), use quantized_index. Components that cannot count structurally raise EnumerationError.

Parameters:

max_fine_bucket (int)

structural_fine_bucket(value, quantizer)[source]

Dominant weighted-component structural bucket (mirrors the plus-of-scaled-children index).

Return type:

int

is_canonical_copy(value, coarse_bin, quantizer)[source]

Stateless dedup: keep value only at its dominant (best-weighted) component’s bin.

The canonical bin is the coarse bin of the minimum, over components, of the component’s structural fine bucket shifted by the weight term. O(K) model evaluations, no state.

Parameters:

coarse_bin (int)

Return type:

bool

class MixtureEnumerator(dist)[source]

Bases: DistributionEnumerator

Enumerator over the deduplicated union of weighted component supports.

Parameters:

dist (MixtureDistribution)

class MixtureSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler that draws a latent component and then samples from that component.

Parameters:

  • dist (MixtureDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw iid samples from a mixture distribution.

The data type drawn from ‘comp_samplers’ is type T, corresponding to the data type of the mixture components.

If size is None, a single sample (of data type T) is drawn and returned. If size is not None, ‘size’-iid mixture samples are drawn and returned as a List with data type List[T].

With batched=True (default) each component sampler is invoked once with the number of draws assigned to it and the results are scattered back into draw order. Because every component sampler owns an independent RandomState, this yields the same draws as the legacy per-draw loop (batched=False) but far faster.

Parameters:

  • size (Optional[int]) – Number of iid samples to draw.

  • batched (bool) – Vectorize component draws (default); set False for the legacy per-draw loop.

Returns:

Data type T or List[T].

Return type:

list[Any] | Any

class MixtureAccumulator(accumulators, keys=(None, None), name=None, init='kmeans++')[source]

Bases: SequenceEncodableStatisticAccumulator

EM accumulator for mixture weights and component sufficient statistics.

Parameters:

  • accumulators (Sequence[SequenceEncodableStatisticAccumulator])

  • keys (tuple[str | None, str | None])

  • name (str | None)

  • init (str)

seq_update(x, weights, estimate)[source]

Accumulate a vectorized EM E-step from encoded observations.

Responsibilities are computed from estimate using the same log-sum-exp normalization as MixtureDistribution.seq_posterior. Rows where every component is impossible receive zero responsibility, so an observation with zero model probability cannot fabricate component sufficient statistics.

Parameters:

  • x (T1) – Encoded observation batch.

  • weights (ndarray) – Non-negative observation weights.

  • estimate (MixtureDistribution) – Previous EM iterate used to compute responsibilities.

Return type:

None

update(x, weight, estimate)[source]

Accumulate one weighted raw observation under an EM estimate.

The observation is routed to each component accumulator with weight * estimate.posterior(x)[k].

Parameters:

  • x (T) – Raw observation.

  • weight (float) – Observation weight.

  • estimate (MixtureDistribution) – Previous EM iterate used to compute responsibilities.

Return type:

None

initialize(x, weight, rng)[source]

Initialize component sufficient statistics from one observation.

This path draws a responsibility vector from a Dirichlet distribution and delegates responsibility-weighted initialization to every component accumulator. It does so for every init strategy: k-means++ seeding needs the whole batch to place its centers, and this method sees one observation at a time. The scalar and vectorized initializations therefore agree only under init="dirichlet"; use seq_initialize() to get the configured strategy.

Parameters:

  • x (T) – Raw observation.

  • weight (float) – Observation weight.

  • rng (RandomState) – Random state used to seed component initializers.

Return type:

None

seq_initialize(x, weights, rng)[source]

Initialize component sufficient statistics from encoded observations.

With init="kmeans++" the method uses a numeric feature matrix when one can be extracted from the encoded data. Ragged, object, hetero, or non-finite encodings fall back to Dirichlet responsibilities rather than mutating input data or forcing an invalid numeric representation.

Parameters:

  • x (T1) – Encoded observation batch.

  • weights (ndarray) – Non-negative observation weights.

  • rng (RandomState) – Random state used to seed component initializers.

Return type:

None

combine(suff_stat)[source]

Merge serialized mixture sufficient statistics into this accumulator.

Parameters:

suff_stat (tuple[ndarray, tuple[T2, ...]]) – (component_counts, component_suff_stats) tuple.

Returns:

self for accumulator chaining.

Return type:

MixtureAccumulator

value()[source]

Return serialized mixture sufficient statistics.

Returns:

(component_counts, component_suff_stats) where the second item contains one serialized child accumulator value per component.

Return type:

tuple[ndarray, tuple[Any, …]]

from_value(x)[source]

Restore this accumulator from serialized sufficient statistics.

Parameters:

x (tuple[ndarray, tuple[T2, ...]]) – (component_counts, component_suff_stats) tuple.

Returns:

self after restoring child accumulator state.

Return type:

MixtureAccumulator

scale(c)[source]

Scale component counts and delegate child sufficient statistics.

Parameters:

c (float)

Return type:

MixtureAccumulator

key_merge(stats_dict)[source]

Merge keyed mixture statistics into a shared statistics dictionary.

Parameters:

stats_dict (dict[str, Any]) – Mutable shared sufficient-statistics mapping keyed by estimator key names.

Return type:

None

key_replace(stats_dict)[source]

Replace local keyed statistics from a shared statistics dictionary.

Parameters:

stats_dict (dict[str, Any]) – Shared sufficient-statistics mapping keyed by estimator key names.

Return type:

None

acc_to_encoder()[source]

Return an encoder assembled from the component accumulators.

Return type:

MixtureDataEncoder

class MixtureAccumulatorFactory(factories, keys=(None, None), name=None, init='kmeans++')[source]

Bases: StatisticAccumulatorFactory

Factory for mixture accumulators built from component accumulator factories.

Parameters:

  • factories (Sequence[StatisticAccumulatorFactory])

  • keys (tuple[str | None, str | None])

  • name (str | None)

  • init (str)

make()[source]

Return a fresh mixture accumulator with fresh component accumulators.

Return type:

MixtureAccumulator

class MixtureEstimator(estimators, fixed_weights=None, suff_stat=None, pseudo_count=None, name=None, keys=(None, None), prior=None, w_min=0.0, robust=False, init=None)[source]

Bases: ParameterEstimator

Estimator for mixture weights and component distributions from EM sufficient statistics.

Parameters:

  • estimators (Sequence[ParameterEstimator])

  • fixed_weights (list[float] | np.ndarray | None)

  • suff_stat (np.ndarray | None)

  • pseudo_count (float | None)

  • name (str | None)

  • keys (tuple[str | None, str | None])

  • prior (SequenceEncodableProbabilityDistribution | None)

  • w_min (float)

  • robust (bool)

  • init (str | None)

accumulator_factory()[source]

Return a mixture accumulator factory matching the component estimators.

Return type:

MixtureAccumulatorFactory

get_prior()[source]

Return the joint mixture prior, or None for a plain MLE estimator.

When a weight prior is attached the joint prior is the (weight_prior, tuple(component priors)) pair produced by mixture_prior().

Return type:

SequenceEncodableProbabilityDistribution | None

set_prior(prior)[source]

Attach a weight prior (and optional per-component priors).

With a (symmetric) Dirichlet weight prior the estimator switches to the conjugate MAP weight update; component priors, when supplied, are delegated to each component estimator via estimator.set_prior (those carry out their own conjugate updates). prior=None leaves the estimator a plain MLE estimator (byte-identical behaviour).

Parameters:

prior (SequenceEncodableProbabilityDistribution | None)

Return type:

None

model_log_density(model)[source]

Log density of the model parameters under this estimator’s prior (ELBO global term).

Returns the Dirichlet weight-prior log-density evaluated at model.w plus the sum of each component estimator’s model_log_density at the corresponding component model. Returns 0.0 for a plain MLE estimator with no priors anywhere.

Parameters:

model (MixtureDistribution)

Return type:

float

estimate(nobs, suff_stat)[source]

Estimate a mixture distribution from EM sufficient statistics.

suff_stat is (component_counts, component_suff_stats). Component parameters are delegated to the child estimators. Mixture weights follow the fixed-weight, conjugate-prior, pseudo-count, or plain-MLE path selected by the estimator configuration. Plain-MLE weights may be floored by w_min and are always renormalized.

Parameters:

  • nobs (float | None) – Unused compatibility argument from ParameterEstimator.

  • suff_stat (tuple[ndarray, tuple[Any, ...]]) – Serialized mixture sufficient statistics.

Returns:

Fitted MixtureDistribution.

Return type:

MixtureDistribution

class MixtureDataEncoder(encoder)[source]

Bases: DataSequenceEncoder

Encoder for homogeneous or heterogeneous mixture component encodings.

Parameters:

encoder (DataSequenceEncoder | Sequence[DataSequenceEncoder])

encoding_signature()[source]

Column-layout identity of this mixture encoder: its kind plus each child’s signature.

Composes so that a leaf which distinguishes “same encoder” from “same encoding” (see BinomialDataEncoder.encoding_signature) is not overruled by a wrapper falling back to __eq__. A child that does not implement it contributes itself, so equality still decides for that branch.

Return type:

tuple

seq_encode(x)[source]

Sequence encode a sequence of iid observations drawn from the mixture distribution.

For a homogeneous mixture this delegates to the single shared component encoder. For a heterogeneous mixture each component encoder encodes the data separately and the encodings are bundled in a _HeteroMixtureEncoded wrapper.

Parameters:

x (Sequence[T]) – A Sequence of iid observations drawn from a mixture distribution with component distributions consistent with the per-component encoders.

Returns:

Encoded sequence (single shared encoding, or a per-component wrapper).

Return type:

Any

row_count(x)[source]

Return the common observation count across shared or heterogeneous encodings.

Parameters:

x (Any)

Return type:

int

class MixtureFisherView(dist)[source]

Bases: FixedFisherView

Complete-data Fisher view for finite mixture distributions.

Coordinates are component assignment indicators followed by each component’s sufficient statistics gated by that assignment. Observed data map to posterior-expected complete-data statistics.

Parameters:

dist (Any)

structured_statistics(x, estimate=None, weight=1.0)[source]

Return mixture responsibility statistics and weighted component statistics for one observation.

Parameters:
Return type:

Any

class ProbabilisticCircuitStatistics(sum_counts, leaf_statistics, leaf_counts, observation_mass, schema_version=1)[source]

Bases: Sequence[Any]

Mass-aware circuit statistics with backward two-slot iteration.

Parameters:
leaf(scope, dist)[source]

A leaf node: an existing mixle dist over the variable indices scope (an int or a tuple).

Parameters:
Return type:

_Node

prod(children)[source]

A product node over children with PAIRWISE-DISJOINT scopes (the decomposability requirement).

Parameters:

children (list[_Node])

Return type:

_Node

summ(children, w=None)[source]

A sum (mixture) node over children that share ONE scope (smoothness); w are mixing weights.

Parameters:
Return type:

_Node

class ProbabilisticCircuitDistribution(root, num_vars, lns_step=None)[source]

Bases: SequenceEncodableProbabilityDistribution

A sum-product network density; build with leaf()/prod()/summ() then pass the root.

Parameters:

  • root (_Node)

  • num_vars (int)

  • lns_step (float | None)

log_density(x)[source]

Return the log-density of one full observation by an upward circuit pass.

Parameters:

x (Any)

Return type:

float

seq_log_density(x)[source]

Return vectorized log-densities for encoded observations.

Parameters:

x (Any)

Return type:

ndarray

dist_to_encoder()[source]

Return the encoder that projects observations into each leaf scope.

Return type:

ProbabilisticCircuitEncoder

sampler(seed=None)[source]

Return an ancestral sampler for this fixed circuit.

Parameters:

seed (int | None)

Return type:

ProbabilisticCircuitSampler

estimator(pseudo_count=None)[source]

Return an EM estimator for this fixed circuit structure.

Parameters:

pseudo_count (float | None)

Return type:

Any

with_params(new_nodes, new_leaf_dists)[source]

A new circuit with the same structure but re-estimated sum-weights / leaf parameters (M-step output).

Parameters:
Return type:

ProbabilisticCircuitDistribution

class ProbabilisticCircuitEncoder(leaf_dists, leaf_scope)[source]

Bases: DataSequenceEncoder

Encode each leaf’s projected columns once with the leaf’s own encoder (shared across EM iterations).

Parameters:
seq_encode(x)[source]

Encode a batch for every leaf distribution using its projected scope.

Parameters:

x (Any)

Return type:

dict[int, Any]

row_count(x)[source]

Return the common number of encoded circuit rows.

Parameters:

x (Any)

Return type:

int

class ProbabilisticCircuitSampler(dist, seed=None)[source]

Bases: DistributionSampler

Ancestral top-down sampling: a sum draws one child by its weights, a product recurses into all.

Parameters:

  • dist (ProbabilisticCircuitDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw one observation or size iid observations from the circuit.

Parameters:
Return type:

Any

class ProbabilisticCircuitAccumulator(nodes, leaf_scope, leaf_estimators)[source]

Bases: SequenceEncodableStatisticAccumulator

E-step sufficient statistics: per-sum-node expected child counts + per-leaf weighted statistics.

The E-step is the circuit FLOW (Peharz et al. EM-for-SPNs): an upward forward gives each node’s log-value, a downward pass gives each node’s log-context lc = derivative of the root log-density w.r.t. that node (the posterior the node is active). A sum node’s per-child responsibility is then exp(lc[sum] + value[child] + log_w - value[sum]) (its expected count), and a leaf’s responsibility is exp(lc[leaf]) (the weight for its sufficient statistic).

Parameters:
seq_update(enc, weights, estimate)[source]

Update circuit-flow responsibilities and leaf sufficient statistics.

Parameters:

  • enc (dict[int, Any])

  • weights (Any)

  • estimate (ProbabilisticCircuitDistribution)

Return type:

None

update(x, weight, estimate)[source]

Update from one weighted observation.

Parameters:

  • x (Any)

  • weight (float)

  • estimate (ProbabilisticCircuitDistribution)

Return type:

None

initialize(x, weight, rng)[source]

Initialize sum counts and leaf statistics from one weighted observation.

Parameters:
Return type:

None

seq_initialize(enc, weights, rng)[source]

Initialize sum counts and leaf statistics from encoded observations.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge sum-node expected counts and leaf accumulator values.

Parameters:

suff_stat (Any)

Return type:

ProbabilisticCircuitAccumulator

value()[source]

Return sum-node expected counts and leaf sufficient statistics.

Return type:

ProbabilisticCircuitStatistics

from_value(x)[source]

Restore sum-node and leaf sufficient statistics from value output.

Parameters:

x (Any)

Return type:

ProbabilisticCircuitAccumulator

scale(c)[source]

Scale sum-node and leaf sufficient statistics by a constant.

Parameters:

c (float)

Return type:

ProbabilisticCircuitAccumulator

key_merge(stats_dict)[source]

Delegate keyed merges to the leaf accumulators.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)[source]

Delegate keyed replacements to the leaf accumulators.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()[source]

Return an encoder based on the current leaf accumulator estimates.

Return type:

ProbabilisticCircuitEncoder

class ProbabilisticCircuitAccumulatorFactory(nodes, leaf_scope, leaf_estimators)[source]

Bases: StatisticAccumulatorFactory

Create accumulators for fixed-structure probabilistic-circuit EM.

Parameters:
make()[source]

Create an empty probabilistic-circuit accumulator.

Return type:

ProbabilisticCircuitAccumulator

class ProbabilisticCircuitEstimator(dist, pseudo_count=None)[source]

Bases: ParameterEstimator

Fits a fixed-structure circuit by EM: renormalize each sum node’s weights to its expected child counts, and re-estimate each leaf from its responsibility-weighted sufficient statistic.

Parameters:

  • dist (ProbabilisticCircuitDistribution)

  • pseudo_count (float | None)

accumulator_factory()[source]

Return a factory for circuit-flow sufficient-statistic accumulators.

Return type:

ProbabilisticCircuitAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate sum-node weights and leaf distributions from accumulated circuit flows.

Parameters:
Return type:

ProbabilisticCircuitDistribution

class ProbabilisticPCADistribution(w, mu, sigma2, name=None, keys=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Probabilistic PCA: x ~ N(mu, W W^T + sigma2 I) with q latent factors.

Parameters:
classmethod compute_capabilities()[source]

Return compute-backend metadata for PPCA scoring.

property w: ndarray

Return an owned snapshot of the factor loadings.

property mu: ndarray

Return an owned snapshot of the mean.

property sigma2: float

Return the immutable isotropic noise variance.

property inv_covar: ndarray

Return an owned snapshot of the cached covariance inverse.

transform(x)[source]

Return the posterior mean of the latent factors E[z | x] = M^{-1} W^T (x - mu).

Parameters:

x (Sequence[float] | ndarray)

Return type:

ndarray

density(x)[source]

Return the probability density at a single observation.

Parameters:

x (Sequence[float] | ndarray)

Return type:

float

log_density(x)[source]

Return the log-density at a single observation.

Parameters:

x (Sequence[float] | ndarray)

Return type:

float

seq_log_density(x)[source]

Return vectorized log-density values for sequence-encoded observations.

Parameters:

x (ndarray)

Return type:

ndarray

backend_seq_log_density(x, engine)[source]

Engine-neutral vectorized log-density for encoded data.

Parameters:
Return type:

Any

sampler(seed=None)[source]

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

ProbabilisticPCASampler

estimator(pseudo_count=None)[source]

Return a closed-form ML estimator with the latent dimension fixed at this model’s q.

Parameters:

pseudo_count (float | None)

Return type:

ProbabilisticPCAEstimator

dist_to_encoder()[source]

Return the data encoder used by this distribution for vectorized methods.

Return type:

ProbabilisticPCADataEncoder

class ProbabilisticPCASampler(dist, seed=None)[source]

Bases: DistributionSampler

Draw iid observations x = mu + W z + sigma * eps from a PPCA model.

Parameters:

  • dist (ProbabilisticPCADistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw size iid vectors (shape (d,) when size is None, else (size, d)).

Parameters:
Return type:

ndarray

class ProbabilisticPCAAccumulator(dim=None, keys=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulate the weighted count, mean, and second-moment matrix (the PPCA sufficient statistics).

Parameters:

  • dim (int | None)

  • keys (str | None)

update(x, weight, estimate)[source]

Accumulate weighted count, sum, and second moment for one vector.

Parameters:

  • x (ndarray)

  • weight (float)

  • estimate (ProbabilisticPCADistribution | None)

Return type:

None

initialize(x, weight, rng)[source]

Initialize the sufficient statistics with one weighted vector.

Parameters:
Return type:

None

seq_update(x, weights, estimate)[source]

Accumulate weighted count, sum, and second moment for encoded vectors.

Parameters:

  • x (ndarray)

  • weights (ndarray)

  • estimate (ProbabilisticPCADistribution | None)

Return type:

None

seq_initialize(x, weights, rng)[source]

Initialize the sufficient statistics from encoded vectors.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge serialized PPCA sufficient statistics into this accumulator.

Parameters:

suff_stat (tuple[float, ndarray | None, ndarray | None])

Return type:

ProbabilisticPCAAccumulator

value()[source]

Return the total weight, weighted sum, and weighted second moment.

Return type:

tuple[float, ndarray | None, ndarray | None]

from_value(x)[source]

Restore the accumulator from serialized PPCA sufficient statistics.

Parameters:

x (tuple[float, ndarray | None, ndarray | None])

Return type:

ProbabilisticPCAAccumulator

acc_to_encoder()[source]

Return an encoder compatible with PPCA vector observations.

Return type:

ProbabilisticPCADataEncoder

class ProbabilisticPCAAccumulatorFactory(dim=None, keys=None)[source]

Bases: StatisticAccumulatorFactory

Factory for ProbabilisticPCAAccumulator.

Parameters:

  • dim (int | None)

  • keys (str | None)

make()[source]

Create an empty PPCA accumulator.

Return type:

ProbabilisticPCAAccumulator

class ProbabilisticPCAEstimator(latent_dim, dim=None, min_sigma2=_MIN_SIGMA2, name=None, keys=None)[source]

Bases: ParameterEstimator

Closed-form maximum-likelihood estimator for PPCA (Tipping & Bishop eigen-solution).

Parameters:

  • latent_dim (int)

  • dim (int | None)

  • min_sigma2 (float)

  • name (str | None)

  • keys (str | None)

accumulator_factory()[source]

Return a factory for PPCA sufficient-statistic accumulators.

Return type:

ProbabilisticPCAAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate PPCA parameters from weighted first and second moments.

Parameters:
Return type:

ProbabilisticPCADistribution

class ProbabilisticPCADataEncoder(dim=None)[source]

Bases: DataSequenceEncoder

Encode a sequence of length-d real vectors into an (n, d) float array.

Parameters:

dim (int | None)

seq_encode(x)[source]

Validate and encode observations as a two-dimensional float array.

Parameters:

x (Sequence[Sequence[float]] | ndarray)

Return type:

ndarray

class QuantizedHMMFitDiagnostics(converged, fixed_theta, theta, objective, iterations_per_start, converged_starts, termination_reasons, selected_start, selected_start_converged, scalar_optimizer_evaluations, split_count, schema_version=1, algorithm='multistart_coordinate_ascent')[source]

Bases: object

Immutable convergence and provenance record for one quantized M-step.

Parameters:

  • converged (bool)

  • fixed_theta (bool)

  • theta (float)

  • objective (float)

  • iterations_per_start (tuple[int, ...])

  • converged_starts (tuple[bool, ...])

  • termination_reasons (tuple[str, ...])

  • selected_start (int)

  • selected_start_converged (bool)

  • scalar_optimizer_evaluations (int)

  • split_count (int)

  • schema_version (int)

  • algorithm (str)

exception QuantizedHMMOptimizationError(message, diagnostics)[source]

Bases: RuntimeError

Raised when the quantized coordinate or scalar optimization fails.

Parameters:

  • message (str)

  • diagnostics (QuantizedHMMFitDiagnostics)

Return type:

None

class QuantizedHiddenMarkovModelDistribution(theta, levels, transition_exponents, emission_exponents, initial_exponents=None, init_mode='quantized', k_max=None, len_dist=NullDistribution(), name=None, terminal_values=None, use_numba=None, terminal_states=None, fit_diagnostics=None)[source]

Bases: HiddenMarkovModelDistribution

Hidden Markov model distribution with quantized observation summaries.

Parameters:

  • theta (float)

  • levels (Sequence[Any])

  • transition_exponents (Sequence[Sequence[int]] | np.ndarray)

  • emission_exponents (Sequence[Sequence[int]] | np.ndarray)

  • initial_exponents (Sequence[int] | np.ndarray | None)

  • init_mode (str)

  • k_max (int | None)

  • len_dist (SequenceEncodableProbabilityDistribution | None)

  • name (str | None)

  • terminal_values (set | None)

  • use_numba (bool | None)

  • terminal_states (set[int] | Sequence[int] | None)

  • fit_diagnostics (QuantizedHMMFitDiagnostics | None)

compute_declaration()[source]

Return the generated-compute declaration for the quantized HMM.

classmethod left_to_right(theta, levels, transition_exponents, emission_exponents, initial_exponents=None, **kwargs)[source]

Construct a left-to-right (upper-triangular) quantized HMM.

transition_exponents must be upper triangular: every entry strictly below the diagonal is a structural zero (negative exponent), so the hidden-state path is monotone non-decreasing (a Bakis chain). This makes a sentence’s state paths exactly its monotone segmentations – only polynomially many in the length (O(L^{n-1})) rather than the n^L of a general HMM – which bounds the path/sequence ambiguity. When the per-state emission supports are additionally disjoint the model is unambiguous (one path per sentence); then the structural descending-probability seek/unrank coincides with the exact marginal order (up to quantization granularity, no path over-count), which a general HMM’s structural seek cannot.

Raises:

ValueError – if transition_exponents is not square or not upper triangular.

Parameters:
Return type:

QuantizedHiddenMarkovModelDistribution

to_fisher(**kwargs)[source]

Forward-backward Fisher view for the quantized HMM.

estimator(pseudo_count=None)[source]

Return an estimator matching this quantized HMM configuration.

Parameters:

pseudo_count (Optional[float]) – Per-cell pseudo count for the initial, transition, and emission expected counts. When None, unobserved cells become structural zeros.

Returns:

QuantizedHiddenMarkovEstimator object.

Return type:

QuantizedHiddenMarkovEstimator

enumerator()[source]

Return an exact descending-probability enumerator over observation sequences.

For the ordinary (length-distribution) case this is the quantized-HMM-specialized enumerator, which avoids constructing per-state categorical streams and uses the cached quantized emission log-probability matrix directly. For the terminal_values stopping-time case it delegates to HiddenMarkovModelEnumerator, which implements that support (the quantized specialization only covers the length-distribution path).

determinize(max_states=1 << 16)[source]

Weighted determinization (Mohri 1997; Mohri & Riley 2002) of this terminal-value quantized HMM into a DeterminizedSequenceDistribution.

Rebuilds the machine over belief states (exact rational arithmetic) so each sequence has a single path and edge weights multiply to the exact marginal – yielding exact, duplicate-free n-best sequences (not n-best paths). Requires terminal_values. Raises EnumerationError if the belief expansion exceeds max_states (the twins property fails – not finitely determinizable; keep the original HMM’s exact O(index) enumerate-and-bin path instead).

Parameters:

max_states (int)

class QuantizedHiddenMarkovModelEnumerator(dist)[source]

Bases: DistributionEnumerator

Exact best-first enumerator specialized for QuantizedHiddenMarkovModelDistribution.

Parameters:

dist (QuantizedHiddenMarkovModelDistribution)

class QuantizedHiddenMarkovEstimator(num_states, levels=None, pseudo_count=None, k_max=None, fixed_theta=None, init_mode='quantized', len_estimator=NullEstimator(), name=None, keys=(None, None, None), use_numba=None, max_quant_its=50, split_collapsed=False, split_nats=math.log(2.0))[source]

Bases: ParameterEstimator

Estimate quantized hidden Markov models from accumulated sufficient statistics.

Parameters:

  • num_states (int)

  • levels (Sequence[Any] | None)

  • pseudo_count (float | None)

  • k_max (int | None)

  • fixed_theta (float | None)

  • init_mode (str)

  • len_estimator (ParameterEstimator | None)

  • name (str | None)

  • keys (tuple[str | None, str | None, str | None] | None)

  • use_numba (bool | None)

  • max_quant_its (int)

  • split_collapsed (bool)

  • split_nats (float)

accumulator_factory()[source]

Returns a HiddenMarkovAccumulatorFactory with categorical emission accumulators.

Return type:

HiddenMarkovAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate a QuantizedHiddenMarkovModelDistribution from Baum-Welch expected counts.

Sufficient statistics in arg ‘suff_stat’ are the HiddenMarkovAccumulator value:

suff_stat[0] (int): Number of hidden states. suff_stat[1] (np.ndarray): Initial state counts. suff_stat[2] (np.ndarray): State counts. suff_stat[3] (np.ndarray): State transition counts. suff_stat[4] (Sequence[Dict[Any, float]]): Per-state categorical emission counts. suff_stat[5] (Optional[Any]): Optional sufficient statistics of the length distribution.

Parameters:
Returns:

QuantizedHiddenMarkovModelDistribution object.

Return type:

QuantizedHiddenMarkovModelDistribution

QuantizedHiddenMarkovModelEstimator

alias of QuantizedHiddenMarkovEstimator

class ResponsibilityAttentionDistribution(key_means, emission, position_prior=None, sigma2=1.0, name=None)[source]

Bases: SequenceEncodableProbabilityDistribution

One generative-gate attention head; a mixture over context positions (EM-able).

Parameters:

  • key_means (np.ndarray)

  • emission (np.ndarray)

  • position_prior (np.ndarray | None)

  • sigma2 (float)

  • name (str | None)

density(x)[source]

Density p(query, target | context) at one observation.

Parameters:

x (tuple[Any, Any, int])

Return type:

float

log_density(x)[source]

Log conditional density log p(query, target | context) at one observation (ctx, y, t).

Parameters:

x (tuple[Any, Any, int])

Return type:

float

seq_log_density(x)[source]

Vectorized log p(query, target | context) over an encoded batch -> (n,).

Parameters:

x (tuple[ndarray, ndarray, ndarray])

Return type:

ndarray

predict_proba(context, query)[source]

Predictive target distribution p(target | context, query) (target marginalized over z).

Attention here does not see the target: a_i ∝ pi_i N(query; K[c_i]), then p(target) = sum_i a_i emission[c_i, :]. Accepts a single example or a batch.

Returns:

(T,) for a single example, or (n, T) for a batch.

Parameters:
Return type:

ndarray

sampler(seed=None)[source]

Return a sampler for synthetic responsibility-attention observations.

Parameters:

seed (int | None)

Return type:

ResponsibilityAttentionSampler

estimator(pseudo_count=None)[source]

Return a closed-form EM estimator for this attention head.

Parameters:

pseudo_count (float | None)

Return type:

ResponsibilityAttentionEstimator

dist_to_encoder()[source]

Return the encoder for contexts, query vectors, and targets.

Return type:

ResponsibilityAttentionDataEncoder

class ResponsibilityAttentionSampler(dist, seed=None)[source]

Bases: DistributionSampler

Joint generative sampler (adds a uniform-distinct context prior so it is a full generator).

Parameters:

  • dist (ResponsibilityAttentionDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw one observation or size iid synthetic observations.

Parameters:
Return type:

Any

class ResponsibilityAttentionAccumulator(num_symbols, context_length, query_dim, num_targets, keys=None, name=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulates the additive sufficient statistics for the responsibility-attention M-step.

Parameters:

  • num_symbols (int)

  • context_length (int)

  • query_dim (int)

  • num_targets (int)

  • keys (str | None)

  • name (str | None)

update(x, weight, estimate)[source]

Update sufficient statistics from one weighted observation.

Parameters:

  • weight (float)

  • estimate (ResponsibilityAttentionDistribution | None)

Return type:

None

seq_update(x, weights, estimate)[source]

Update sufficient statistics from encoded observations.

Parameters:

  • weights (ndarray)

  • estimate (ResponsibilityAttentionDistribution | None)

Return type:

None

initialize(x, weight, rng)[source]

Initialize sufficient statistics from one weighted observation.

Parameters:
Return type:

None

seq_initialize(x, weights, rng)[source]

Initialize sufficient statistics with random attention responsibilities.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge key, emission, position, and query-scatter statistics.

Return type:

ResponsibilityAttentionAccumulator

value()[source]

Return a snapshot of additive attention sufficient statistics.

from_value(x)[source]

Restore attention sufficient statistics from value output.

Return type:

ResponsibilityAttentionAccumulator

acc_to_encoder()[source]

Return the encoder compatible with this accumulator.

Return type:

ResponsibilityAttentionDataEncoder

class ResponsibilityAttentionAccumulatorFactory(num_symbols, context_length, query_dim, num_targets, keys=None, name=None)[source]

Bases: StatisticAccumulatorFactory

Create accumulators for responsibility-attention EM statistics.

Parameters:

  • num_symbols (int)

  • context_length (int)

  • query_dim (int)

  • num_targets (int)

  • keys (str | None)

  • name (str | None)

make()[source]

Create an empty responsibility-attention accumulator.

Return type:

ResponsibilityAttentionAccumulator

class ResponsibilityAttentionEstimator(num_symbols, context_length, query_dim, num_targets, *, sigma2=1.0, estimate_sigma2=False, min_sigma2=1e-6, emission_smoothing=1e-6, pseudo_count=None, name=None, keys=None)[source]

Bases: ParameterEstimator

Estimate a ResponsibilityAttentionDistribution by closed-form EM M-steps.

Parameters:

  • num_symbols (int)

  • context_length (int)

  • query_dim (int)

  • num_targets (int)

  • sigma2 (float)

  • estimate_sigma2 (bool)

  • min_sigma2 (float)

  • emission_smoothing (float)

  • pseudo_count (float | None)

  • name (str | None)

  • keys (str | None)

accumulator_factory()[source]

Return a factory for responsibility-attention sufficient-statistic accumulators.

Return type:

ResponsibilityAttentionAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate key means, emissions, position prior, and optional gate variance.

Parameters:

nobs (float | None)

Return type:

ResponsibilityAttentionDistribution

class ResponsibilityAttentionDataEncoder(num_symbols=None, context_length=None, query_dim=None, num_targets=None)[source]

Bases: DataSequenceEncoder

Encodes a sequence of (context, query, target) triples into stacked arrays.

seq_encode(x)[source]

Encode (context, query, target) observations.

Parameters:

x (Sequence[tuple[Any, Any, int]])

Return type:

tuple[ndarray, ndarray, ndarray]

row_count(x)[source]

Return the number of observations represented by x.

The default covers simple array/list encodings and structurally aligned tuples/dicts. Encoders with flattened, ragged, or otherwise specialized layouts must override this method so metadata wrappers can verify counts instead of trusting a caller assertion.

Return type:

int

sequence_to_triples(tokens, context_length, *, num_symbols=None, embeddings=None)[source]

Unroll a token sequence into the (context, query, target) triples this head consumes.

For each prediction point p (with a full window behind it and a next token ahead), emits context = tokens[p-N+1 : p+1] (the N most recent tokens, including the current one), query derived from the current token tokens[p] (its embedding, or a one-hot), and target = tokens[p+1]. This is the bridge from a real sequence to the iid-triple leaf.

Scope: with the single-hop leaf this yields an attention-weighted next-token model whose predictive is a function of the current token (an attention-flavoured bigram) – because the emission is keyed by the attended token’s identity, the prediction depends only on which token the query selects. In-context copy / induction (example-specific values) is what the multi-hop / stackable extension adds; this helper is the plumbing both share.

Parameters:

  • tokens (Sequence[int]) – a 1-D sequence of integer token ids.

  • context_length (int) – the attention window N.

  • num_symbols (int | None) – vocabulary size (for the one-hot query dim); inferred from tokens if omitted.

  • embeddings (ndarray | None) – optional (S, D) embedding table; if given the query is embeddings[token], otherwise a one-hot of dimension num_symbols.

Returns:

A list of (context, query, target) triples ready for mixle.inference.optimize().

Return type:

list[tuple[ndarray, ndarray, int]]

class PhaseSchedule[source]

Bases: object

Maps a position t in a length-L sequence to a phase index in [0, n_phases).

phase(t, length)[source]

Return the phase index for position t in a sequence of length.

Parameters:
Return type:

int

to_dict()[source]

Serialize the schedule to a JSON-compatible dictionary.

Return type:

dict[str, Any]

static from_dict(d)[source]

Deserialize a schedule produced by to_dict().

Parameters:

d (dict[str, Any])

Return type:

PhaseSchedule

class Homogeneous[source]

Bases: PhaseSchedule

One phase for everything – the ordinary time-homogeneous HMM.

phase(t, length)[source]

Return the single homogeneous phase.

Parameters:
Return type:

int

to_dict()[source]

Serialize the homogeneous schedule.

Return type:

dict[str, Any]

class ByPosition(cap)[source]

Bases: PhaseSchedule

Absolute position: phi(t, L) = min(t, cap - 1) (positions past cap-1 share the last phase).

Parameters:

cap (int)

phase(t, length)[source]

Return the absolute-position phase capped at the final phase.

Parameters:
Return type:

int

to_dict()[source]

Serialize the absolute-position schedule.

Return type:

dict[str, Any]

class ByRelativePosition(bins)[source]

Bases: PhaseSchedule

Relative position: phi(t, L) = min(bins - 1, floor(bins * t / L)) – progress through the sequence.

Parameters:

bins (int)

phase(t, length)[source]

Return the relative-position phase for t / length.

Parameters:
Return type:

int

to_dict()[source]

Serialize the relative-position schedule.

Return type:

dict[str, Any]

class ByLength(boundaries)[source]

Bases: PhaseSchedule

Length-conditional: phase is the bucket of L against sorted boundaries (constant within a seq).

With boundaries = [5, 10] there are three phases: L <= 5, 5 < L <= 10, L > 10.

Parameters:

boundaries (Sequence[int])

phase(t, length)[source]

Return the length-bucket phase for the sequence length.

Parameters:
Return type:

int

to_dict()[source]

Serialize the length-bucket schedule.

Return type:

dict[str, Any]

class ScheduledHiddenMarkovModelDistribution(inits, transitions, emissions, schedule, len_dist=None, name=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Phase-indexed (length-/position-conditional) HMM. See the module docstring for the modeling story.

Parameters:

  • inits (np.ndarray)

  • transitions (np.ndarray)

  • emissions (list[list[Any]])

  • schedule (PhaseSchedule)

  • len_dist (Any)

  • name (str | None)

log_density(x)[source]

Return the log likelihood of one scheduled HMM sequence.

Parameters:

x (list[Any])

Return type:

float

seq_log_density(x)[source]

Score a batch of scheduled HMM sequences.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)[source]

Return a sampler for scheduled HMM sequences.

Parameters:

seed (int | None)

Return type:

ScheduledHMMSampler

estimator(pseudo_count=None)[source]

Create the matching estimator from this distribution’s own components.

Every phase/state cell retains its own estimator, so heterogeneous emission families survive fitting. The length estimator mirrors len_dist when one is modeled.

Parameters:

pseudo_count (float | None)

Return type:

ScheduledHMMEstimator

dist_to_encoder()[source]

Return the pass-through scheduled HMM encoder.

Return type:

ScheduledHMMDataEncoder

class ScheduledHMMSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for scheduled HMM sequences.

Parameters:

  • dist (ScheduledHiddenMarkovModelDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw one sequence or a list of sequences.

Parameters:
Return type:

Any

class ScheduledHMMDataEncoder[source]

Bases: DataSequenceEncoder

Pass-through encoder for scheduled HMM sequence observations.

seq_encode(x)[source]

Encode scheduled HMM records as sequence lists.

Parameters:

x (list[list[Any]])

Return type:

list[list[Any]]

row_count(x)[source]

Return the number of pass-through scheduled-HMM records.

Parameters:

x (Any)

Return type:

int

class ScheduledHMMAccumulator(n_states, schedule, emission_factories, len_factory=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulator for phase-pooled scheduled HMM EM sufficient statistics.

Parameters:

  • n_states (int)

  • schedule (PhaseSchedule)

  • emission_factories (Any)

  • len_factory (Any)

update(x, weight, estimate)[source]

Accumulate sufficient statistics from one weighted sequence.

Parameters:

  • x (list[Any])

  • weight (float)

  • estimate (ScheduledHiddenMarkovModelDistribution)

Return type:

None

seq_update(x, weights, estimate)[source]

Accumulate weighted sufficient statistics from a batch.

Parameters:

  • x (Any)

  • weights (ndarray)

  • estimate (ScheduledHiddenMarkovModelDistribution)

Return type:

None

initialize(x, weight, rng)[source]

Initialize sufficient statistics with random soft state responsibilities.

Parameters:
Return type:

None

seq_initialize(x, weights, rng)[source]

Initialize sufficient statistics from a weighted batch.

Parameters:
Return type:

None

combine(other)[source]

Merge serialized scheduled HMM sufficient statistics.

Parameters:

other (Any)

Return type:

ScheduledHMMAccumulator

value()[source]

Return serialized scheduled HMM sufficient statistics.

Return type:

tuple

from_value(value)[source]

Restore accumulator state from serialized sufficient statistics.

Parameters:

value (tuple)

Return type:

ScheduledHMMAccumulator

acc_to_encoder()[source]

Return the encoder associated with this accumulator.

Return type:

ScheduledHMMDataEncoder

class ScheduledHMMAccumulatorFactory(n_states, schedule, emission_estimators, len_estimator=None)[source]

Bases: StatisticAccumulatorFactory

Factory for scheduled HMM accumulators.

Parameters:

  • n_states (int)

  • schedule (PhaseSchedule)

  • emission_estimators (Any)

  • len_estimator (Any)

make()[source]

Create a fresh scheduled HMM accumulator.

Return type:

ScheduledHMMAccumulator

class ScheduledHMMEstimator(n_states, schedule, emission_estimator, len_estimator=None, pseudo_count=1e-8, name=None, *, emission_estimators=None)[source]

Bases: ParameterEstimator

EM estimator for a ScheduledHiddenMarkovModelDistribution with a fixed schedule.

emission_estimator is the backward-compatible homogeneous estimator prototype. emission_estimators may instead provide an exact phase-by-state grid, preserving heterogeneous families.

Parameters:

  • n_states (int)

  • schedule (PhaseSchedule)

  • emission_estimator (Any)

  • len_estimator (Any)

  • pseudo_count (float)

  • name (str | None)

  • emission_estimators (Sequence[Sequence[Any]] | None)

accumulator_factory()[source]

Return the accumulator factory used by this estimator.

Return type:

ScheduledHMMAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate phase-indexed initial, transition, emission, and length models.

Parameters:
Return type:

ScheduledHiddenMarkovModelDistribution

SegmentalHiddenMarkovDistribution

alias of SegmentalHiddenMarkovModelDistribution

class SegmentalHiddenMarkovModelDistribution(emissions, w=MISSING, transitions=MISSING, len_dist=NullDistribution(), name=None, weights=MISSING, terminal_states=None)[source]

Bases: SequenceEncodableProbabilityDistribution

HMM whose states emit arbitrary segment-valued distributions.

Observations are lists of segment objects. For example, with SequenceDistribution(GaussianDistribution(...), len_dist=...) as an emission, each state emits a variable-length list of real values.

Parameters:
compute_capabilities()[source]

Describe backend support shared by emissions and optional length model.

compute_declaration()[source]

Return a composite compute declaration for the segmental HMM.

property topics: list[SequenceEncodableProbabilityDistribution]

Compatibility alias with HiddenMarkovModelDistribution terminology.

density(x)[source]

Return the probability density or mass at a single observation.

Parameters:

x (Sequence[Any])

Return type:

float

log_density(x)[source]

Return the log-density or log-mass at a single observation.

Parameters:

x (Sequence[Any])

Return type:

float

seq_log_density(x)[source]

Return vectorized log-density values for sequence-encoded observations.

Parameters:

x (tuple[ndarray, ndarray, tuple[Any, ...], Any | None])

Return type:

ndarray

backend_seq_log_density(x, engine)[source]

Engine-neutral segmental-HMM forward scoring for encoded segment sequences.

Parameters:
Return type:

Any

sampler(seed=None)[source]

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

SegmentalHiddenMarkovSampler

estimator(pseudo_count=None)[source]

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

SegmentalHiddenMarkovEstimator

dist_to_encoder()[source]

Return the data encoder used by this distribution for vectorized methods.

Return type:

SegmentalHiddenMarkovDataEncoder

enumerator()[source]

Enumerate segment sequences in descending marginal probability order.

The segmental HMM has the standard HMM forward semantics – each position emits one segment from its state’s distribution, scored independently – so it reuses HiddenMarkovModelEnumerator directly via its per-state emission (topics), log_w, log_transitions, and len_dist. Each segment is drawn from the union of the per-state emission supports, so every emission distribution must itself support enumeration (and a length distribution must be modeled).

Return type:

DistributionEnumerator

class SegmentalHiddenMarkovSampler(dist, seed=None)[source]

Bases: DistributionSampler

Draw iid segmental-HMM observations.

Parameters:

  • dist (SegmentalHiddenMarkovModelDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw one segment sequence, or size iid segment sequences.

Parameters:
Return type:

list[Any] | list[list[Any]]

class SegmentalHiddenMarkovAccumulator(accumulators, len_accumulator=NullAccumulator(), keys=(None, None, None), name=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Baum-Welch accumulator for segmental HMMs.

Parameters:

  • accumulators (Sequence[SequenceEncodableStatisticAccumulator])

  • len_accumulator (SequenceEncodableStatisticAccumulator | None)

  • keys (tuple[str | None, str | None, str | None] | None)

  • name (str | None)

update(x, weight, estimate)[source]

Update sufficient statistics for one observed segment sequence.

Parameters:

  • x (Sequence[Any])

  • weight (float)

  • estimate (SegmentalHiddenMarkovModelDistribution)

Return type:

None

initialize(x, weight, rng)[source]

Randomly initialize sufficient statistics for one segment sequence.

Parameters:
Return type:

None

seq_initialize(x, weights, rng)[source]

Randomly initialize state and emission statistics for encoded sequences.

Parameters:
Return type:

None

seq_update(x, weights, estimate)[source]

Update encoded-sequence statistics with Baum-Welch posteriors.

Parameters:
Return type:

None

seq_update_engine(x, weights, estimate, engine)[source]

Engine-resident Baum-Welch E-step (numpy or torch).

Reuses hmm_engine_forward_backward over the sequence-contiguous segment encoding, producing the same initial/state/transition/emission statistics as the host seq_update.

combine(suff_stat)[source]

Merge another segmental-HMM sufficient-statistic value.

Parameters:

suff_stat (tuple[int, ndarray, ndarray, ndarray, Sequence[Any], Any | None])

Return type:

SegmentalHiddenMarkovAccumulator

value()[source]

Return transition, emission, and sequence-length sufficient statistics.

Return type:

tuple[int, ndarray, ndarray, ndarray, tuple[Any, …], Any | None]

from_value(x)[source]

Replace this accumulator from serialized sufficient statistics.

Parameters:

x (tuple[int, ndarray, ndarray, ndarray, Sequence[Any], Any | None])

Return type:

SegmentalHiddenMarkovAccumulator

scale(c)[source]

Scale all weight-linear sufficient statistics by c.

Parameters:

c (float)

Return type:

SegmentalHiddenMarkovAccumulator

key_merge(stats_dict)[source]

Merge keyed initial, transition, emission, and length statistics.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

key_replace(stats_dict)[source]

Replace keyed initial, transition, emission, and length statistics.

Parameters:

stats_dict (dict[str, Any])

Return type:

None

acc_to_encoder()[source]

Return an encoder compatible with the emission and length accumulators.

Return type:

SegmentalHiddenMarkovDataEncoder

class SegmentalHiddenMarkovAccumulatorFactory(factories, len_factory=NullAccumulatorFactory(), keys=(None, None, None), name=None)[source]

Bases: StatisticAccumulatorFactory

Factory for SegmentalHiddenMarkovAccumulator.

Parameters:

  • factories (Sequence[StatisticAccumulatorFactory])

  • len_factory (StatisticAccumulatorFactory | None)

  • keys (tuple[str | None, str | None, str | None] | None)

  • name (str | None)

make()[source]

Create a fresh segmental-HMM accumulator.

Return type:

SegmentalHiddenMarkovAccumulator

class SegmentalHiddenMarkovEstimator(estimators, len_estimator=NullEstimator(), pseudo_count=(None, None), name=None, keys=(None, None, None), terminal_states=None)[source]

Bases: ParameterEstimator

Baum-Welch estimator for SegmentalHiddenMarkovModelDistribution.

Parameters:
accumulator_factory()[source]

Return an accumulator factory for Baum-Welch sufficient statistics.

Return type:

SegmentalHiddenMarkovAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate initial, transition, emission, and length distributions.

Parameters:
Return type:

SegmentalHiddenMarkovModelDistribution

class SegmentalHiddenMarkovDataEncoder(emission_encoders, len_encoder=NullDataEncoder())[source]

Bases: DataSequenceEncoder

Encode a batch of segment sequences for a segmental HMM.

Parameters:

  • emission_encoders (Sequence[DataSequenceEncoder])

  • len_encoder (DataSequenceEncoder | None)

seq_encode(x)[source]

Flatten segment sequences and encode them for every state emission.

Parameters:

x (Sequence[Sequence[Any]])

Return type:

tuple[ndarray, ndarray, tuple[Any, …], Any | None]

row_count(x)[source]

Return and validate the number of encoded segment sequences.

Parameters:

x (Any)

Return type:

int

SegmentalHiddenMarkovModelSampler

alias of SegmentalHiddenMarkovSampler

SegmentalHiddenMarkovModelAccumulator

alias of SegmentalHiddenMarkovAccumulator

SegmentalHiddenMarkovModelAccumulatorFactory

alias of SegmentalHiddenMarkovAccumulatorFactory

SegmentalHiddenMarkovModelEstimator

alias of SegmentalHiddenMarkovEstimator

SegmentalHiddenMarkovModelDataEncoder

alias of SegmentalHiddenMarkovDataEncoder

class SemiSupervisedHiddenMarkovModelDistribution(topics, transitions, len_dist=None, name=None, keys=None, use_numba=None, terminal_states=None)[source]

Bases: SequenceEncodableProbabilityDistribution

HMM with shared emissions/transitions where each observation may carry a per-position state prior.

compute_capabilities()[source]

Declare the legacy NumPy execution path for semi-supervised HMM inference.

density(x)[source]

Return the probability of one semi-supervised HMM observation.

Return type:

float

log_density(x)[source]

Return the log-likelihood of one (emissions, state_prior) observation.

Return type:

float

seq_log_density(x)[source]

Return vectorized log-likelihoods for encoded semi-supervised HMM observations.

Return type:

ndarray

density_semantics()[source]

Return the joined density semantics of emission and optional length distributions.

sampler(seed=None)[source]

Return a sampler for emission sequences with no external state priors.

estimator(pseudo_count=None)[source]

Return a Baum-Welch estimator for transitions, emissions, and optional length.

dist_to_encoder()[source]

Return the encoder for emission sequences, priors, and optional lengths.

class SemiSupervisedHiddenMarkovSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sample emission sequences from the HMM with a uniform initial state distribution.

Priors are external side information, so sampled observations carry None as their prior.

Parameters:

dist (SemiSupervisedHiddenMarkovModelDistribution)

sample(size=None, *, batched=True)[source]

Draw one observation or a list of observations with None state priors.

Parameters:

batched (bool)

SemiSupervisedHiddenMarkovAccumulator

alias of SemiSupervisedHiddenMarkovEstimatorAccumulator

SemiSupervisedHiddenMarkovAccumulatorFactory

alias of SemiSupervisedHiddenMarkovEstimatorAccumulatorFactory

class SemiSupervisedHiddenMarkovEstimator(estimators, len_estimator=None, pseudo_count=None, name=None, keys=(None, None), terminal_states=None)[source]

Bases: ParameterEstimator

Estimate transitions, emissions, and optional length from semi-supervised HMM statistics.

accumulator_factory()[source]

Return a factory for semi-supervised HMM sufficient-statistic accumulators.

estimate(nobs, suff_stat)[source]

Estimate the HMM transition matrix and child distributions from accumulated statistics.

class SemiSupervisedHiddenMarkovDataEncoder(emission_encoder, len_encoder=None)[source]

Bases: DataSequenceEncoder

Encode a sequence of (emission_seq, state_prior) observations for the semi-supervised HMM.

seq_encode(x)[source]

Encode observations as emission lists, priors, optional lengths, and raw lengths.

row_count(x)[source]

Return and validate the number of encoded semi-supervised sequences.

Return type:

int

SemiSupervisedHiddenMarkovModelSampler

alias of SemiSupervisedHiddenMarkovSampler

SemiSupervisedHiddenMarkovModelAccumulator

alias of SemiSupervisedHiddenMarkovEstimatorAccumulator

SemiSupervisedHiddenMarkovModelAccumulatorFactory

alias of SemiSupervisedHiddenMarkovEstimatorAccumulatorFactory

SemiSupervisedHiddenMarkovModelEstimator

alias of SemiSupervisedHiddenMarkovEstimator

SemiSupervisedHiddenMarkovModelDataEncoder

alias of SemiSupervisedHiddenMarkovDataEncoder

class SemiSupervisedMixtureDistribution(components, w=MISSING, name=None, weights=MISSING)[source]

Bases: SequenceEncodableProbabilityDistribution

SemiSupervisedMixtureDistribution models observations (value, prior) where the optional prior labels re-weight the mixture weights over the listed components.

Parameters:
compute_capabilities()[source]

Declare generated-compute support inherited from mixture components.

compute_declaration()[source]

Return the generated-compute declaration for the semi-supervised mixture.

density(x)[source]

Density of the semi-supervised mixture at observation x.

See log_density() for details.

Parameters:

x (Tuple[T0, Optional[Sequence[Tuple[int, T1]]]]) – Observation (value, prior).

Returns:

Density at x.

Return type:

float

log_density(x)[source]

Log-density of the semi-supervised mixture at observation x = (value, prior).

If prior is None this is the standard mixture log-density. Otherwise the mixture weights are restricted to the components listed in the prior, re-weighted by the prior probabilities, and re-normalized before mixing the component log-densities.

Parameters:

x (Tuple[T0, Optional[Sequence[Tuple[int, T1]]]]) – Observation (value, prior), where prior is an optional sequence of (component index, probability) pairs.

Returns:

Log-density at x.

Return type:

float

posterior(x)[source]

Posterior probability of each component for observation x = (value, prior).

Components not listed in the prior (when a prior is present) receive posterior 0.

Parameters:

x (Tuple[T0, Optional[Sequence[Tuple[int, T1]]]]) – Observation (value, prior).

Returns:

Numpy array of length num_components containing the component posteriors.

Return type:

ndarray

seq_log_density(x)[source]

Vectorized evaluation of the log-density on sequence encoded data x.

Parameters:

x (E) – Sequence encoded data produced by SemiSupervisedMixtureDataEncoder.seq_encode().

Returns:

Numpy array of log-densities, one entry per encoded observation.

Return type:

ndarray

backend_seq_log_density(x, engine)[source]

Engine-neutral semi-supervised mixture log-density for encoded observations.

Parameters:
Return type:

Any

seq_posterior(x)[source]

Vectorized component posteriors on sequence encoded data x.

Parameters:

x (E) – Sequence encoded data produced by SemiSupervisedMixtureDataEncoder.seq_encode().

Returns:

Numpy array of shape (number of observations, num_components) of posteriors.

Return type:

ndarray

density_semantics()[source]

Return exact-or-approximate density semantics joined from component models.

sampler(seed=None)[source]

Creates a SemiSupervisedMixtureSampler for sampling component values.

Parameters:

seed (Optional[int]) – Seed for the random number generator used in sampling.

Returns:

SemiSupervisedMixtureSampler object.

Return type:

SemiSupervisedMixtureSampler

estimator(pseudo_count=None)[source]

Creates a SemiSupervisedMixtureEstimator with one child estimator per component.

Parameters:

pseudo_count (Optional[float]) – Used to inflate the sufficient statistics of the mixture weights.

Returns:

SemiSupervisedMixtureEstimator object.

Return type:

SemiSupervisedMixtureEstimator

dist_to_encoder()[source]

Creates a SemiSupervisedMixtureDataEncoder for encoding sequences of (value, prior) observations.

Returns:

SemiSupervisedMixtureDataEncoder object.

Return type:

SemiSupervisedMixtureDataEncoder

enumerator()[source]

Enumeration is not well-defined for semi-supervised mixtures.

Observations pair a component value with exogenous prior labels: the model defines no distribution over the prior part, so the support over (value, prior) pairs cannot be enumerated with consistent probabilities.

Raises:

EnumerationError always. –

Return type:

DistributionEnumerator

class SemiSupervisedMixtureSampler(dist, seed=None)[source]

Bases: DistributionSampler

SemiSupervisedMixtureSampler draws component values from a SemiSupervisedMixtureDistribution.

Parameters:

  • dist (SemiSupervisedMixtureDistribution)

  • seed (int | None)

sample(size=None, *, batched=True)[source]

Draw ‘size’ component values from the mixture (no prior labels are generated).

Parameters:

  • size (Optional[int]) – Number of independent samples. If None a single value is returned.

  • batched (bool)

Returns:

A single component value if size is None, else a list of ‘size’ component values.

Return type:

Sequence[Any] | Any

SemiSupervisedMixtureAccumulator

alias of SemiSupervisedMixtureEstimatorAccumulator

SemiSupervisedMixtureAccumulatorFactory

alias of SemiSupervisedMixtureEstimatorAccumulatorFactory

class SemiSupervisedMixtureEstimator(estimators, suff_stat=None, pseudo_count=None, keys=(None, None), name=None)[source]

Bases: ParameterEstimator

SemiSupervisedMixtureEstimator estimates a SemiSupervisedMixtureDistribution from aggregated sufficient statistics.

Parameters:
accumulator_factory()[source]

Create an accumulator factory from the child estimators.

Returns:

Factory configured for this estimator.

Return type:

SemiSupervisedMixtureEstimatorAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate a SemiSupervisedMixtureDistribution from aggregated sufficient statistics.

The mixture weights are the normalized component counts, optionally regularized by pseudo_count and the stored suff_stat weights.

Parameters:

  • nobs (Optional[float]) – Not used. Kept for consistency.

  • suff_stat (Tuple[np.ndarray, Tuple[SS0, ...]]) – Component counts and component sufficient statistics, as returned by SemiSupervisedMixtureEstimatorAccumulator.value().

Returns:

Estimated distribution.

Return type:

SemiSupervisedMixtureDistribution

class SemiSupervisedMixtureDataEncoder(encoder, num_components=None)[source]

Bases: DataSequenceEncoder

SemiSupervisedMixtureDataEncoder encodes sequences of (value, prior) observations using a shared component encoder for the values and flat arrays for the prior labels.

Parameters:

  • encoder (DataSequenceEncoder)

  • num_components (int | None)

row_count(x)[source]

Return the observation count recorded in the encoded payload.

The base implementation infers a count only from payloads whose leading axis is unambiguously the observation axis. This encoder’s payload is a heterogeneous 4-tuple – count, encoded values, prior arrays, raw observations – whose members have differing and partly ragged lengths (the prior arrays are indexed by labelled row, so they are shorter than the data whenever any observation is unlabelled). Inference declines that, so without this override every metadata wrapper that verifies a row count raised NotImplementedError against a perfectly ordinary semi-supervised fit.

Element 0 is the count, written by seq_encode() from the input length itself.

Parameters:

x (Any)

Return type:

int

seq_encode(x)[source]

Encode a sequence of iid (value, prior) observations for vectorized seq_ calls.

The encoding is a tuple of length 4:

rv[0] (int): Number of observations.
rv[1]: The values encoded by the shared component encoder.
rv[2]: Prior arrays ((row index, component index, prob, log prob), per-row prior
    sums, per-row has-prior flags).
rv[3]: The raw observations (used by seq_initialize).
Parameters:

x (Sequence[Tuple[T0, Optional[Sequence[Tuple[int, T1]]]]]) – Observations.

Returns:

See description above.

Return type:

tuple[int, Any, tuple[tuple[ndarray, ndarray, ndarray, ndarray], ndarray, ndarray], Sequence[tuple[T0, Sequence[tuple[int, T1]] | None]]]

class SparseScore(lower, upper, exact, n_scored)[source]

Bases: object

A certified bracket on a mixture’s log p(x) from scoring only the top components.

lower <= log p(x) <= upper always holds; exact is True when the bracket collapsed (all components scored, or max_components covered the support). n_scored components were exactly evaluated. When exact, lower == upper == log p(x).

Parameters:
log_density_sup(dist)[source]

Maximum log p(x) of a single distribution, or None if unbounded / unknown.

A finite value certifies the component can contribute at most this much, which is what makes the dropped-tail bound sound. None (unbounded density or an unrecognized family) forces the exact path.

Parameters:

dist (Any)

Return type:

float | None

sparse_mixture_score(mixture, x, max_components)[source]

Score mixture.log_density(x) exactly on the top max_components and certify the rest.

Ranks components by the low-cost upper bound log w_k + sup_k (x-independent), scores the exact contribution of only the top ones, and bounds the dropped tail by the remaining upper bounds. Returns a SparseScore bracket. If any positive-weight component’s density is unbounded (log_density_sup is None), falls back to exact full scoring (lower == upper, no speedup).

Parameters:
Return type:

SparseScore

collapse_identical(mixture)[source]

Merge components that are identical (same family + parameters) by summing their weights – EXACT.

A fitted huge mixture often carries duplicate components; pooling exact duplicates leaves log p(x) unchanged while shrinking K. Identity requires the same fully qualified family and canonical serialized state. A component that cannot produce that representation is conservatively left separate. ‘Blend a mixture of closed forms’ with zero approximation.

Parameters:

mixture (Any)

Return type:

Any

collapse_gaussian_mixture(mixture, max_components)[source]

Reduce an all-Gaussian mixture to <= max_components by moment-matching the nearest pair.

Greedily merges the two closest components (by mean) into the single Gaussian preserving their combined weight, mean, and variance, until max_components remain. Approximate (it widens), but it preserves the OVERALL mixture mean and variance exactly – the analytic ‘collapse a huge mixture of closed forms’ that lets scoring/sampling scale.

Parameters:

  • mixture (Any)

  • max_components (int)

Return type:

Any

class TransitionOperator[source]

Bases: object

A row-stochastic state-transition operator behind the HMM forward-backward.

Subclasses provide the two linear maps the recursions need plus an M-step from expected transition mass. forward/backward must be consistent with as_matrix (forward(a) == a @ A, backward(v) == A @ v); the low-overhead operators never materialize A.

to_dict()[source]

Return a safe JSON-compatible representation of this operator.

Return type:

dict[str, Any]

classmethod from_dict(payload)[source]

Reconstruct an operator from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

TransitionOperator

forward(alpha)[source]

Push a state-belief row vector one transition forward.

Parameters:

alpha (ndarray)

Return type:

ndarray

backward(v)[source]

Pull an emission-weighted belief vector one transition backward.

Parameters:

v (ndarray)

Return type:

ndarray

as_matrix()[source]

Materialize the transition matrix when an explicit matrix is needed.

Return type:

ndarray

new_accumulator()[source]

Create an empty transition sufficient-statistic accumulator.

Return type:

Any

accumulate(acc, alpha_t, w_next, scale)[source]

Add one transition’s expected mass. alpha_t is the (normalized) forward belief at t, w_next = b_{t+1} * beta_{t+1} the emission-weighted backward belief at t+1, scale the forward normalizer c_{t+1} (so the per-step posterior transition mass is exact).

Parameters:
Return type:

None

estimate(acc)[source]

Estimate a transition operator of the same structural family from accumulated statistics.

Parameters:

acc (Any)

Return type:

TransitionOperator

random_accumulator(rng)[source]

A randomly-filled accumulator whose estimate yields a random (structured) transition – used to seed EM when there is no warm start. Fills new_accumulator shapes (nested) with noise.

Return type:

Any

class HMMFitDiagnostics(algorithm, converged, iterations, termination_reason, log_likelihood_trace, initial_log_likelihood, final_log_likelihood, final_absolute_delta, final_relative_delta, monotone, n_sequences, total_weight, approximate)[source]

Bases: object

Machine-readable receipt for a structured-family EM fit.

Parameters:

  • algorithm (str)

  • converged (bool)

  • iterations (int)

  • termination_reason (str)

  • log_likelihood_trace (tuple[float, ...])

  • initial_log_likelihood (float)

  • final_log_likelihood (float)

  • final_absolute_delta (float | None)

  • final_relative_delta (float | None)

  • monotone (bool)

  • n_sequences (int)

  • total_weight (float)

  • approximate (bool)

class HMMFitResult(model, trace, diagnostics)[source]

Bases: tuple

Backward-compatible (model, trace) result with a fit diagnostics receipt.

Parameters:

  • model (Any)

  • trace (list[float])

  • diagnostics (HMMFitDiagnostics)

property model

Return the fitted model.

property log_likelihood_trace

Return the accepted model log-likelihood trajectory.

class DenseTransition(a, prior=None)[source]

Bases: TransitionOperator

The usual dense K x K row-stochastic transition (O(K^2) forward-backward).

prior (a K x K pseudocount matrix) is added to the expected counts before each M-step re-normalization – a Dirichlet/MAP transition. A diagonal prior is a sticky self-transition bias (see sticky_transition()); a flat prior is symmetric-Dirichlet smoothing.

Parameters:

  • a (np.ndarray)

  • prior (np.ndarray | None)

forward(alpha)[source]

Push a state-belief row vector forward with the dense matrix.

backward(v)[source]

Pull a vector backward with the dense matrix.

as_matrix()[source]

Return the dense row-stochastic transition matrix.

new_accumulator()[source]

Create dense expected-transition-count storage.

accumulate(acc, alpha_t, w_next, scale)[source]

Accumulate one dense expected-transition-count contribution.

estimate(acc)[source]

Estimate a row-normalized dense transition from expected counts.

class LowRankTransition(g, phi)[source]

Bases: TransitionOperator

A = G @ Phi: each state’s next-state distribution is a mix of r shared transition profiles.

G is K x r row-stochastic (state -> profile mixing), Phi is r x K row-stochastic (profile -> next-state). A = G @ Phi is K x K row-stochastic with rank <= r. Forward ((alpha @ G) @ Phi), backward (G @ (Phi @ v)) and the M-step are all O(K r) – never forming A – and the parameter count is 2 K r instead of K^2.

Parameters:

  • g (np.ndarray)

  • phi (np.ndarray)

forward(alpha)[source]

Push a state-belief row vector through the low-rank transition.

backward(v)[source]

Pull a vector backward through the low-rank transition.

as_matrix()[source]

Materialize the implied dense transition matrix.

new_accumulator()[source]

Create state-profile and profile-next sufficient-statistic storage.

accumulate(acc, alpha_t, w_next, scale)[source]

Accumulate one exact low-rank transition contribution.

estimate(acc)[source]

Estimate low-rank transition factors from accumulated statistics.

class SparseTransition(n_states, edges, values=None)[source]

Bases: TransitionOperator

Only the given (from, to) edges are allowed (left-to-right / banded HMMs). Forward, backward and the M-step are O(#edges) – transitions outside the edge set stay exactly zero through EM, so the structure is preserved. Build edges yourself or with left_to_right_edges() / banded_edges().

Parameters:

n_states (int)

forward(alpha)[source]

Push a state-belief row vector through the sparse transition.

backward(v)[source]

Pull a vector backward through the sparse transition.

as_matrix()[source]

Materialize the sparse transition as a dense matrix.

new_accumulator()[source]

Create one expected-count slot per allowed edge.

accumulate(acc, alpha_t, w_next, scale)[source]

Accumulate one sparse expected-transition contribution.

estimate(acc)[source]

Estimate a sparse transition over the same allowed edge set.

class BlockDiagonalTransition(blocks)[source]

Bases: TransitionOperator

Independent sub-chains: the states partition into blocks and transitions stay within a block.

A model whose initial state picks a block and then evolves inside it – a mixture of regimes that do not switch. Build it from any sub-operators (each block can itself be dense or low-rank). Exact, block-local forward-backward and M-step.

forward(alpha)[source]

Push each block’s belief mass forward independently.

backward(v)[source]

Pull each block’s vector backward independently.

as_matrix()[source]

Materialize the block-diagonal transition matrix.

new_accumulator()[source]

Create one transition accumulator per block.

accumulate(acc, alpha_t, w_next, scale)[source]

Accumulate block-local expected transition mass.

estimate(acc)[source]

Estimate each block and return a new block-diagonal transition.

class KroneckerTransition(op1, op2)[source]

Bases: TransitionOperator

Factorial HMM: the state is the pair (s1, s2) of two chains evolving in parallel, with A = A1 (x) A2 (Kronecker). State index is i1 * K2 + i2.

Forward-backward uses the reshape identity (alpha @ (A1 (x) A2) reshapes to A1^T @ M @ A2), so a step is O(K1 K2 (K1 + K2)) instead of O((K1 K2)^2) – the whole point of a factorial HMM. The E-step is exact over the joint state; the M-step is the standard factorial marginal update (each factor re-estimated from the marginalized joint transition mass), verified to keep EM monotone.

Parameters:

  • op1 (TransitionOperator)

  • op2 (TransitionOperator)

forward(alpha)[source]

Push a joint belief through the Kronecker-factorized transition.

backward(v)[source]

Pull a joint vector backward through the Kronecker transition.

as_matrix()[source]

Materialize the dense Kronecker transition matrix.

new_accumulator()[source]

Create marginal transition-count accumulators for both factors.

accumulate(acc, alpha_t, w_next, scale)[source]

Accumulate exact marginalized factor transition statistics.

estimate(acc)[source]

Estimate both Kronecker factors from marginalized expected counts.

class FinalStateEnumeration(hmm, len_dist)[source]

Bases: object

Result of _final_state_enumerate(): top_k(k) -> [(sequence, log_prob), …] descending.

top_k(k)[source]

Return the top k final-state-constrained sequences.

class StructuredHMM(emissions, pi, transition, emission_estimators=None, keys=(None, None), name=None, len_dist=None, terminal_states=None, final_states=None)[source]

Bases: object

An HMM whose transition is a TransitionOperator (dense / low-rank / a combinator).

emissions is one observation distribution per state; pi the initial-state distribution; transition any TransitionOperator. The scaled forward-backward and EM call the operator’s forward/backward/accumulate/estimate, so a low-rank or factorial transition runs the SAME inference at its own cost (O(K r) for low-rank). emission_estimators (one per state) drives the emission M-step; default reuses emissions[k].estimator().

Parameters:

transition (TransitionOperator)

viterbi(seq)[source]

Most-likely state path (Viterbi / max-product). Uses the transition matrix, so it works for any operator; O(T K^2) – a read-out, not the EM hot loop.

posterior_decode(seq)[source]

Per-position MAP state argmax_k P(z_t = k | x) from the forward-backward posteriors gamma.

enumerator()[source]

Enumerate observation sequences in descending marginal probability (top_k / rank / seek / nucleus / certified estimates). Enumeration depends only on pi, the transition MATRIX, the emissions and a length distribution – not on the operator’s internal structure – so it reuses the built-in HMM enumerator (an A*-style best-first search over the trellis) on the dense matrix. Requires len_dist (a distribution over sequence length) and enumerable (discrete) emissions.

dist_to_enumerator()[source]

Return the sequence enumerator for this HMM.

state_posteriors(seq)[source]

The full smoothing posteriors gamma[t,k] = P(z_t = k | x).

seq_log_density(seqs)[source]

Score a batch of observation sequences.

Return type:

ndarray

sampler(seed=None)[source]

Return a sampler for observation sequences.

fit(seqs, *, max_its=50, tol=1e-6, fast=True, weights=None)[source]

Fit by Baum-Welch through the transition operator.

fast=True uses the numba-jitted dense forward-backward when possible. The returned HMMFitResult remains unpackable as (fitted_hmm, loglik_trace) and also exposes a validated diagnostics receipt.

Parameters:
class StructuredHMMDataEncoder[source]

Bases: DataSequenceEncoder

Sequences pass through as lists – the structured forward-backward scores raw observations through the per-state emission log_density (no flattened columnar encoding; composability over raw speed).

seq_encode(x)[source]

Encode sequences as lists without changing their observations.

row_count(x)[source]

Return the number of pass-through structured-HMM records.

class StructuredHMMAccumulator(emission_accumulators, transition_proto, keys=(None, None))[source]

Bases: SequenceEncodableStatisticAccumulator

Baum-Welch E-step accumulator: per-sequence forward-backward, accumulating initial-state mass, transition-operator mass, and per-state weighted emission statistics.

update(x, weight, estimate)[source]

Accumulate sufficient statistics from one weighted sequence.

seq_update(x, weights, estimate)[source]

Run forward-backward and accumulate weighted sufficient statistics for a batch.

seq_initialize(x, weights, rng)[source]

Initialize sufficient statistics with random soft state responsibilities.

combine(suff_stat)[source]

Merge serialized HMM sufficient statistics.

value()[source]

Return serialized HMM sufficient statistics.

from_value(x)[source]

Restore accumulator state from serialized sufficient statistics.

scale(factor)[source]

Multiply the running statistics by factor – the decay primitive online/streaming Baum-Welch (StreamingEstimator) uses to fold a new batch into a forgetting running estimate.

acc_to_encoder()[source]

Return the encoder associated with this accumulator.

key_merge(store)[source]

Merge tied initial or transition sufficient statistics into store.

key_replace(store)[source]

Replace tied initial or transition statistics from store.

class StructuredHMMAccumulatorFactory(emission_estimators, transition_proto, keys)[source]

Bases: StatisticAccumulatorFactory

Factory for structured HMM accumulators.

make()[source]

Create a fresh structured HMM accumulator.

class StructuredHMMEstimator(emission_estimators, transition_proto, keys=(None, None), name=None, len_dist=None, terminal_states=None)[source]

Bases: ParameterEstimator

Estimator (M-step) for a StructuredHMM: re-estimates pi, the transition OPERATOR (any structure – dense/low-rank/combinator), and each state’s emission from the Baum-Welch statistics. keys=(init_key, trans_key) tie the initial / transition parameters across HMMs that share them.

accumulator_factory()[source]

Return the accumulator factory used by this estimator.

estimate(nobs, suff_stat)[source]

Estimate an HMM from Baum-Welch sufficient statistics.

class InputOutputHMM(emissions, pi, transitions, emission_estimators=None, name=None, terminal_states=None)[source]

Bases: object

Input-output HMM (IOHMM): an exogenous discrete input u_t selects which transition governs each step. Holds one TransitionOperator per input symbol; the emission is per-state. Data is (obs_seq, input_seq) pairs where input_seq[t] in {0..M-1} drives the transition from t to t+1.

Lets a covariate steer the dynamics – regime switching driven by an observed control, the difference between a plain HMM and a controlled Markov model. (Input-dependent emissions are a natural extension; here emissions depend on state only.)

seq_log_density(x, input_seqs=None)[source]

Per-sequence forward log-likelihood. Two call forms: - seq_log_density(obs_seqs, input_seqs) – the explicit two-list API; or - seq_log_density(records) – one list of (obs, input)-pair sequences (the 5-part contract).

viterbi(seq, inputs=None)[source]

Most-likely state path (Viterbi / max-product), conditioned on the input/control sequence: step t -> t+1 maximizes over the transition inputs[t] selects. Call as viterbi(record) on one (observation, input)-pair sequence or as viterbi(obs_seq, input_seq). Uses the per-input transition matrices, so it works for any operator; O(T K^2) – a read-out, not the EM hot loop.

posterior_decode(seq, inputs=None)[source]

Per-position MAP state argmax_k P(z_t = k | x, u) from the forward-backward posteriors gamma.

state_posteriors(seq, inputs=None)[source]

The full smoothing posteriors gamma[t,k] = P(z_t = k | x, u), conditioned on the inputs.

sampler(seed=None)[source]

Return a sampler for (observation, input)-pair records along a given control sequence.

fit(obs_seqs, input_seqs, *, max_its=50, tol=1e-6, weights=None)[source]

Fit the IOHMM by weighted Baum-Welch and return a diagnostics-bearing result.

Parameters:
log_density(seq)[source]

Return the log likelihood of one (observation, input) sequence.

dist_to_encoder()[source]

Return the pass-through IOHMM sequence encoder.

estimator(pseudo_count=None)[source]

Return the estimator for this IOHMM structure.

to_dict()[source]

Return a safe JSON-compatible representation of this IOHMM (decodes via load_models).

Return type:

dict[str, Any]

class IOHMMDataEncoder[source]

Bases: DataSequenceEncoder

An IOHMM record is one (obs, input) sequence – a list of (observation, input_symbol) pairs.

seq_encode(x)[source]

Encode IOHMM records as lists of (observation, input) pairs.

row_count(x)[source]

Return the number of pass-through input-output HMM records.

class IOHMMAccumulator(emission_accumulators, transition_protos)[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulator for IOHMM Baum-Welch sufficient statistics.

update(x, weight, estimate)[source]

Accumulate sufficient statistics from one weighted IOHMM record.

seq_update(x, weights, estimate)[source]

Accumulate weighted sufficient statistics from a batch of IOHMM records.

seq_initialize(x, weights, rng)[source]

Initialize IOHMM sufficient statistics with random soft responsibilities.

combine(suff_stat)[source]

Merge serialized IOHMM sufficient statistics.

value()[source]

Return serialized IOHMM sufficient statistics.

from_value(x)[source]

Restore accumulator state from serialized IOHMM statistics.

scale(factor)[source]

Scale IOHMM transition, emission, and initial-state statistics.

acc_to_encoder()[source]

Return the encoder associated with this accumulator.

class IOHMMAccumulatorFactory(emission_estimators, transition_protos)[source]

Bases: StatisticAccumulatorFactory

Factory for IOHMM accumulators.

make()[source]

Create a fresh IOHMM accumulator.

class IOHMMEstimator(emission_estimators, transition_protos, name=None)[source]

Bases: ParameterEstimator

Estimator (M-step) for an InputOutputHMM: re-estimates pi, one transition operator per input symbol (from the per-input expected counts), and each state’s emission.

accumulator_factory()[source]

Return the accumulator factory used by this estimator.

estimate(nobs, suff_stat)[source]

Estimate an IOHMM from Baum-Welch sufficient statistics.

class ExplicitDurationHMM(emissions, pi, transition_matrix, durations, max_duration, name=None)[source]

Bases: object

Hidden semi-Markov model (explicit-duration HMM): each state emits for a random duration drawn from a per-state duration distribution, then switches state (the transition matrix has a zero diagonal – dwell time is modeled explicitly, not as a self-loop). This captures non-geometric state durations a plain HMM cannot.

durations is one length-max_duration probability vector per state (over d = 1..max_duration). A sequence has an exogenous fixed observation horizon. Its final latent duration is right-censored when it extends beyond that horizon, exactly matching sample(length). Forward/EM are O(T * K * max_duration).

forward_loglik(seq)[source]

Fixed-horizon log likelihood, marginalizing a possibly right-censored final duration.

fit(seqs, *, max_its=50, tol=1e-6, weights=None)[source]

Fit the explicit-duration HMM by weighted Baum-Welch with a diagnostics receipt.

Parameters:
log_density(seq)[source]

Return the explicit-duration HMM log likelihood for one sequence.

seq_log_density(x)[source]

Score a batch of observation sequences.

dist_to_encoder()[source]

Return the pass-through explicit-duration sequence encoder.

estimator(pseudo_count=None)[source]

Return the estimator for this explicit-duration HMM structure.

to_structured_hmm(len_dist=None)[source]

The HSMM as an EQUIVALENT StructuredHMM via the remaining-duration expansion: K*D sub-states (k, r) = “state k with r steps left in the segment”. The expanded chain emits from state k at every sub-state, decrements deterministically (k,r)->(k,r-1), and at (k,1) switches segment with A[k,k’]*dur[k’](d’). Reading the first T expanded-chain emissions exactly represents the right-censored fixed horizon, so no final-state constraint is applied. This hands the HSMM the full StructuredHMM read-out API and, with len_dist, enumeration. O(K*D) states.

enumerator(len_dist)[source]

Enumerate fixed-horizon observation sequences in descending marginal probability, given a len_dist over total sequence length. Built on the exact right-censored HMM expansion; .top_k(k) -> [(sequence, log_prob), …].

state_posteriors(seq)[source]

Per-position smoothing posteriors gamma[t, j] = P(z_t = j | obs), marginalizing the durations (sum the posterior of every segment that covers position t). Rows sum to 1.

posterior_decode(seq)[source]

Per-position MAP state argmax_j P(z_t = j | obs).

viterbi_segments(seq)[source]

Most-likely segmentation (max-product over the segment lattice): a list of (state, start, observed_duration) segments covering the sequence, O(T K D). The last tuple describes the observed portion of a right-censored latent duration.

sampler(seed=None)[source]

Return a sampler for explicit-duration HMM sequences.

class EDHMMDataEncoder[source]

Bases: DataSequenceEncoder

An ExplicitDurationHMM record is one observation sequence (the durations are latent).

seq_encode(x)[source]

Encode EDHMM records as observation-sequence lists.

row_count(x)[source]

Return the number of pass-through explicit-duration HMM records.

class EDHMMAccumulator(emission_accumulators, k, d)[source]

Bases: SequenceEncodableStatisticAccumulator

E-step accumulator for an explicit-duration HMM: per-sequence segment posteriors -> initial / transition / per-state DURATION counts + emission occupancy statistics.

update(x, weight, estimate)[source]

Accumulate sufficient statistics from one weighted sequence.

seq_update(x, weights, estimate)[source]

Accumulate weighted explicit-duration HMM statistics from a batch.

seq_initialize(x, weights, rng)[source]

Initialize explicit-duration HMM statistics with random soft responsibilities.

combine(suff_stat)[source]

Merge serialized explicit-duration HMM sufficient statistics.

value()[source]

Return serialized explicit-duration HMM sufficient statistics.

from_value(x)[source]

Restore accumulator state from serialized EDHMM statistics.

scale(factor)[source]

Scale explicit-duration HMM statistics.

acc_to_encoder()[source]

Return the encoder associated with this accumulator.

class EDHMMAccumulatorFactory(emission_estimators, k, d)[source]

Bases: StatisticAccumulatorFactory

Factory for explicit-duration HMM accumulators.

make()[source]

Create a fresh explicit-duration HMM accumulator.

class EDHMMEstimator(emission_estimators, k, d, name=None)[source]

Bases: ParameterEstimator

Estimator (M-step) for an ExplicitDurationHMM: re-estimates pi, the zero-diagonal transition, the per-state DURATION distributions, and each state’s emission from the segment-posterior statistics.

accumulator_factory()[source]

Return the accumulator factory used by this estimator.

estimate(nobs, suff_stat)[source]

Estimate an explicit-duration HMM from segment posterior statistics.

sticky_transition(a, kappa)[source]

A dense transition with a STICKY self-transition prior: kappa pseudocounts on the diagonal favor staying in a state (longer dwell times, cleaner segmentation – the sticky-HMM idea).

Parameters:

kappa (float)

Return type:

DenseTransition

dirichlet_transition(a, alpha)[source]

A dense transition with a symmetric Dirichlet(alpha) smoothing prior on every row (MAP).

Parameters:

alpha (float)

Return type:

DenseTransition

kron_initial(pi1, pi2)[source]

Factorized initial distribution pi1 (x) pi2 for a factorial (Kronecker) HMM – the two chains start independently. Matches a KroneckerTransition so the joint initial respects the factors.

Return type:

ndarray

left_to_right_edges(n_states, skip=1)[source]

Edges for a left-to-right (Bakis) HMM: each state may stay or advance up to skip states.

Parameters:
banded_edges(n_states, bandwidth=1)[source]

Edges for a banded transition: state i connects to i-bandwidth .. i+bandwidth (local time-series).

Parameters:

  • n_states (int)

  • bandwidth (int)

chunked_state_posteriors(hmm, seq, *, chunk, overlap)[source]

State posteriors gamma for one long sequence via overlapping chunks, each run INDEPENDENTLY (embarrassingly parallel). The first chunk uses the model’s pi; interior chunks start from the uniform belief and the overlap context lets the chain forget that wrong boundary – so the kept interior matches the exact forward-backward up to an error that decays at the mixing rate in overlap.

Parameters:

  • hmm (StructuredHMM)

  • chunk (int)

  • overlap (int)

Return type:

ndarray

fit_chunked(hmm, seqs, *, chunk, overlap, max_its=50, workers=0, tol=1e-6, weights=None)[source]

Baum-Welch where each long sequence’s forward-backward is split into overlapping chunks run in PARALLEL (the forgetting property bounds the boundary error). workers>0 runs the per-chunk E-steps on a thread pool (NumPy releases the GIL in its array kernels); workers=0 runs them serially. The interior suff-statistics are accumulated exactly as in StructuredHMM.fit(); this only changes how the E-step is computed, trading a small, overlap-controlled approximation for intra-sequence parallelism. The returned HMMFitResult remains unpackable as (fitted_hmm, loglik_trace) and marks the receipt’s objective as approximate.

Parameters:

  • hmm (StructuredHMM)

  • chunk (int)

  • overlap (int)

  • max_its (int)

  • workers (int)

  • tol (float)

stationary_initial(op, *, iters=2000, tol=1e-13)[source]

The transition’s stationary distribution (pi @ A == pi), by power iteration through op.forward – so it is O(K r) for a low-rank op, never forming A. Use it to COUPLE a StructuredHMM’s initial state to its transition (pi = stationary_initial(transition)): the chain starts in its long-run distribution instead of a free, separately-estimated pi. Answers “do the initial states match the transition?” – they can, by construction.

Parameters:

  • op (TransitionOperator)

  • iters (int)

  • tol (float)

Return type:

ndarray

jit_forward_loglik(hmm)[source]

Compile the scaled forward log-likelihood recursion to a single jax.jit XLA program (lax.scan over time). Returns a callable score(seq) -> float: emission log-densities are evaluated on the host (arbitrary emissions), then the forward scan runs jitted on the transition matrix. Works for any operator (uses as_matrix()); the win is large T / K. Requires the JAX optional extra.

Precision follows the caller’s ambient JAX policy rather than mixle’s: with the default jax_enable_x64=False the scan runs in float32, so the returned score agrees with seq_log_density to roughly float32 relative precision, not to float64. Enable x64 in your own process if you need more – mixle does not set it for you (MXR-080-0147 removed the import-time global mutation that used to make this path silently float64).

Parameters:

hmm (StructuredHMM)

class TreeHiddenMarkovModelDistribution(topics, w=MISSING, transitions=MISSING, len_dist=NullDistribution(), terminal_level=10, name=None, use_numba=None, weights=MISSING)[source]

Bases: SequenceEncodableProbabilityDistribution

Hidden Markov model on a rooted tree with emission distributions of type T.

Data type: Sequence[Tuple[Tuple[int, int], T]] (((node_id, parent_id), emission) per node, root parent -1).

Parameters:
compute_capabilities()[source]

Declare generated-compute support from emission and length children.

compute_declaration()[source]

Return the generated-compute declaration for the tree HMM.

density(x)[source]

Density of tree HMM at a single observed tree x.

See log_density() for details.

Parameters:

x (Sequence[Tuple[D, T]]) – Tree as ((node_id, parent_id), emission) tuples (root parent -1).

Returns:

Density at observation x.

Return type:

float

log_density(x)[source]

Log-density of tree HMM at a single observed tree x.

The hidden states are marginalized out with an upward (beta) message passing recursion over the tree. When a non-null length distribution (len_dist) is set, its contribution to the likelihood (the sum over nodes of len_dist.log_density(num_children)) is included; a NullDistribution length contributes nothing.

Parameters:

x (Sequence[Tuple[D, T]]) – Tree as ((node_id, parent_id), emission) tuples (root parent -1).

Returns:

Log-density at observation x.

Return type:

float

seq_log_density(x)[source]

Vectorized evaluation of log-density at sequence encoded input x.

Dispatches to numba kernels or to the pure-numpy level-by-level recursion depending on which encoding (use_numba) the input was created with.

Parameters:

x (E) – Sequence encoded trees from TreeHiddenMarkovDataEncoder.seq_encode().

Returns:

Numpy array of log-density (float), one entry per encoded tree.

Return type:

ndarray

backend_seq_log_density(x, engine)[source]

Engine-neutral tree-HMM scoring for the pure non-numba encoded layout.

Parameters:
Return type:

Any

seq_posterior(x)[source]

Posterior state membership probabilities for each node of each encoded tree.

Parameters:

x (E) – Sequence encoded trees from TreeHiddenMarkovDataEncoder.seq_encode().

Returns:

List with one numpy array per tree, each of shape (num_nodes, num_states), containing the posterior probability of each hidden state at each node.

Return type:

list[ndarray] | None

viterbi(x)[source]

Most likely hidden state assignment for each node of a single observed tree.

Parameters:

x (Sequence[Tuple[D, T]]) – Tree as ((node_id, parent_id), emission) tuples (root parent -1).

Returns:

Numpy array of int states, one per node of the tree.

Return type:

ndarray

seq_viterbi(x)[source]

Vectorized Viterbi state assignments for sequence encoded trees.

Parameters:

x (E) – Sequence encoded trees from TreeHiddenMarkovDataEncoder.seq_encode().

Returns:

List with one numpy array of int states per tree (one state per node).

Return type:

list[ndarray]

density_semantics()[source]

Return exact-or-approximate density semantics joined from child models.

sampler(seed=None)[source]

Create a sampler for this tree hidden Markov distribution.

Parameters:

seed (Optional[int]) – Used to set seed in random sampler.

Returns:

Sampler bound to this distribution.

Return type:

TreeHiddenMarkovSampler

Raises:

ValueError – If len_dist is a NullDistribution (a length distribution with support on the non-negative integers is required for sampling).

enumerator()[source]

Not supported: observations are rooted trees, not chains.

The chain best-first enumerator (HiddenMarkovModelEnumerator) does not apply – the marginal sums over hidden states on a branching tree whose shape is itself governed by the per-node child-count len_dist, so the support is a set of trees rather than sequences. Use sampler() or the exact log_density instead.

Return type:

DistributionEnumerator

estimator(pseudo_count=None)[source]

Create a TreeHiddenMarkovEstimator with estimators for the topics and length distribution.

Parameters:

pseudo_count (Optional[float]) – Used to inflate sufficient statistics of the initial state weights, transition matrix, topics, and length distribution.

Returns:

TreeHiddenMarkovEstimator object.

Return type:

TreeHiddenMarkovEstimator

dist_to_encoder()[source]

Return a data encoder for iid tree HMM observations.

Return type:

TreeHiddenMarkovDataEncoder

class TreeHiddenMarkovSampler(dist, seed=None)[source]

Bases: DistributionSampler

Sampler for the TreeHiddenMarkovModelDistribution. Draws rooted trees of state-emitted observations.

Parameters:

  • dist (TreeHiddenMarkovModelDistribution)

  • seed (int | None)

sample_state(given_state, size=None)[source]

Draw child state(s) from the transition matrix row of a given parent state.

Parameters:

  • given_state (int) – State of the parent node.

  • size (Optional[int]) – Number of child states to draw. If None, a single int is returned.

Returns:

Single state (int) if size is None, else numpy array of size states.

Return type:

int | ndarray

sample_tree(size=None)[source]

Sample rooted tree(s) of ((node_id, parent_id), emission) tuples.

The root state is drawn from the initial weights, each child state from the transition matrix, and the number of children of each node from the length sampler. Sampling stops once the terminal_level of the distribution is reached.

Parameters:

size (Optional[int]) – Number of trees to draw. If None, a single tree is returned.

Returns:

A single tree (list of ((node_id, parent_id), emission) tuples with root parent -1) if size is None, else a list of size trees.

sample(size=None, *, batched=True)[source]

Draw iid tree observations from the tree HMM.

Parameters:

  • size (Optional[int]) – Number of trees to draw. If None, a single tree is returned.

  • batched (bool)

Returns:

A single tree if size is None, else a list of size trees. See sample_tree().

Raises:

RuntimeError – If no length sampler is available for the number of children.

class TreeHiddenMarkovAccumulator(accumulators, len_accumulator=NullAccumulator(), keys=(None, None, None), name=None, use_numba=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulator for the tree HMM. Tracks initial-state, state, and transition counts plus emission stats.

Parameters:

  • accumulators (Sequence[SequenceEncodableStatisticAccumulator])

  • len_accumulator (SequenceEncodableStatisticAccumulator | None)

  • keys (tuple[str | None, str | None, str | None])

  • name (str | None)

  • use_numba (bool | None)

update(x, weight, estimate)[source]

Update sufficient statistics with a single weighted tree observation.

Encodes the tree and delegates to seq_update().

Parameters:

  • x (Sequence[Tuple[D, T]]) – Tree as ((node_id, parent_id), emission) tuples (root parent -1).

  • weight (float) – Weight for observation.

  • estimate (TreeHiddenMarkovModelDistribution) – Previous estimate used for the E-step.

Return type:

None

initialize(x, weight, rng)[source]

Initialize sufficient statistics with a single weighted tree observation.

Parameters:

  • x (Sequence[Tuple[D, T]]) – Tree as ((node_id, parent_id), emission) tuples (root parent -1).

  • weight (float) – Weight for observation.

  • rng (RandomState) – Random number generator used to draw random state assignments.

Return type:

None

seq_initialize(x, weights, rng)[source]

Vectorized initialization of sufficient statistics from sequence encoded trees.

Hidden states are assigned uniformly at random and the initial-state, state, transition, emission, and length statistics are accumulated under that random assignment.

Parameters:

  • x (E) – Sequence encoded trees from TreeHiddenMarkovDataEncoder.seq_encode().

  • weights (np.ndarray) – Weights for each encoded tree.

  • rng (np.random.RandomState) – Random number generator used to draw state assignments.

Return type:

None

seq_update(x, weights, estimate)[source]

Vectorized E-step update of sufficient statistics from sequence encoded trees.

Runs the upward-downward (Baum-Welch style) recursion under the previous estimate to obtain expected initial-state, state membership, and transition counts, and passes the node posteriors as weights to the emission accumulators (and tree weights to the length accumulator).

Parameters:

  • x (E) – Sequence encoded trees from TreeHiddenMarkovDataEncoder.seq_encode().

  • weights (np.ndarray) – Weights for each encoded tree.

  • estimate (TreeHiddenMarkovModelDistribution) – Previous estimate used for the E-step.

Return type:

None

seq_update_engine(x, weights, estimate, engine)[source]

Engine-resident E-step for the pure (non-numba) tree encoding.

Runs the upward-downward recursion entirely with engine ops (numpy or torch): emission scoring, the level-by-level beta/eta upward pass, and the alpha/xi downward pass all live on the active engine. Node posteriors and aggregated initial/state/transition counts are produced on the engine and converted to host arrays only to feed the child accumulators. Mirrors the numpy branch of seq_update.

Parameters:
Return type:

None

combine(suff_stat)[source]

Combine sufficient statistics from another accumulator into this one.

Parameters:

suff_stat (Tuple) – Tuple of number of states, initial-state counts, state counts, transition counts, emission sufficient statistics per state, and length sufficient statistics.

Returns:

Self, with aggregated sufficient statistics.

Return type:

TreeHiddenMarkovAccumulator

value()[source]

Returns sufficient statistics as a Tuple of number of states, initial-state counts, state counts, transition counts, emission sufficient statistics per state, and length statistics.

Return type:

tuple[int, ndarray, ndarray, ndarray, Sequence[Any], Any | None]

from_value(x)[source]

Set sufficient statistics of accumulator from value x.

Parameters:

x (Tuple) – Tuple of number of states, initial-state counts, state counts, transition counts, emission sufficient statistics per state, and length sufficient statistics.

Returns:

Self, with sufficient statistics set to x.

Return type:

TreeHiddenMarkovAccumulator

scale(c)[source]

Scale all accumulated tree-HMM sufficient statistics in place.

Parameters:

c (float)

Return type:

TreeHiddenMarkovAccumulator

key_merge(stats_dict)[source]

Merge keyed sufficient statistics into stats_dict (and recurse into member accumulators).

Parameters:

stats_dict (Dict[str, Any]) – Dictionary mapping keys to shared sufficient statistics.

Return type:

None

key_replace(stats_dict)[source]

Replace keyed sufficient statistics with values from stats_dict (and recurse into members).

Parameters:

stats_dict (Dict[str, Any]) – Dictionary mapping keys to shared sufficient statistics.

Return type:

None

acc_to_encoder()[source]

Return a data encoder built from the emission and length accumulators.

Return type:

TreeHiddenMarkovDataEncoder

class TreeHiddenMarkovAccumulatorFactory(factories, len_factory=NullAccumulatorFactory(), keys=(None, None, None), name=None, use_numba=None)[source]

Bases: StatisticAccumulatorFactory

Factory for tree hidden Markov accumulators.

Parameters:

  • factories (Sequence[StatisticAccumulatorFactory])

  • len_factory (StatisticAccumulatorFactory)

  • keys (tuple[str | None, str | None, str | None] | None)

  • name (str | None)

  • use_numba (bool | None)

make()[source]

Return a new tree hidden Markov accumulator.

Return type:

TreeHiddenMarkovAccumulator

class TreeHiddenMarkovEstimator(estimators, len_estimator=NullEstimator(), pseudo_count=(None, None), name=None, keys=(None, None, None), use_numba=None)[source]

Bases: ParameterEstimator

Estimator for the TreeHiddenMarkovModelDistribution from aggregated sufficient statistics.

Parameters:

  • estimators (list[ParameterEstimator])

  • len_estimator (ParameterEstimator | None)

  • pseudo_count (float | tuple[float | None, float | None] | None)

  • name (str | None)

  • keys (tuple[str | None, str | None, str | None] | None)

  • use_numba (bool | None)

accumulator_factory()[source]

Return an accumulator factory configured from this estimator.

Return type:

TreeHiddenMarkovAccumulatorFactory

estimate(nobs, suff_stat)[source]

Estimate a TreeHiddenMarkovModelDistribution from sufficient statistics (M-step).

Initial state weights and transition rows are normalized counts, optionally smoothed with the pseudo_count pair. A transition row with no observed outgoing transitions – e.g. every training tree is shallow enough that this state is never seen with a child – is filled uniform rather than left at zero when no pseudo-count is given, so that p_level (the per-level marginal state-occupancy vector derived from w and this matrix, used throughout scoring) stays strictly positive at every level instead of going hard-zero and later poisoning log_density() on a deeper out-of-sample tree with a silent NaN (T4-01). Every such row is disclosed through numerical_repairs() on the returned distribution. Topics and the length distribution are estimated from their respective sufficient statistics.

Parameters:

  • nobs (Optional[float]) – Number of observations.

  • suff_stat (Tuple) – Tuple of number of states, initial-state counts, state counts, transition counts, emission sufficient statistics per state, and length sufficient statistics.

Returns:

Estimated distribution.

Return type:

TreeHiddenMarkovModelDistribution

class TreeHiddenMarkovDataEncoder(emission_encoder, len_encoder=NullDataEncoder(), use_numba=None)[source]

Bases: DataSequenceEncoder

Data encoder for sequences of tree HMM observations (flattens trees into level-indexed arrays).

Parameters:

  • emission_encoder (DataSequenceEncoder)

  • len_encoder (DataSequenceEncoder | None)

  • use_numba (bool | None)

row_count(x)[source]

Return the number of trees represented by either encoded layout.

Parameters:

x (Any)

Return type:

int

seq_encode(x)[source]

Encode a sequence of tree observations for vectorized functions.

Returns a pair (numba_encoding, numpy_encoding) with exactly one entry set, depending on use_numba: the numba encoding stores per-tree slice offsets consumed by the numba kernels, while the numpy encoding (see _seq_encode()) batches nodes across trees by level.

Parameters:

x (Sequence[Sequence[Tuple[D, T]]]) – Sequence of trees of ((node_id, parent_id), emission) tuples (root parent -1).

Returns:

Tuple (E) with the numba encoding in slot 0 (numpy slot None), or vice versa.

Return type:

tuple[tuple[ndarray, tuple[int, ndarray, ndarray, ndarray], tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray, ndarray], Any, tuple[ndarray, Any] | None] | None, tuple[int, ndarray, tuple[ndarray, ndarray, ndarray], tuple[ndarray, ndarray, ndarray, ndarray, list[ndarray], list[ndarray], list[ndarray], list[ndarray], ndarray], Any, tuple[ndarray, Any] | None] | None]

TreeHiddenMarkovModelSampler

alias of TreeHiddenMarkovSampler

TreeHiddenMarkovModelAccumulator

alias of TreeHiddenMarkovAccumulator

TreeHiddenMarkovModelAccumulatorFactory

alias of TreeHiddenMarkovAccumulatorFactory

TreeHiddenMarkovModelEstimator

alias of TreeHiddenMarkovEstimator

TreeHiddenMarkovModelDataEncoder

alias of TreeHiddenMarkovDataEncoder

find_level(parents)[source]

Find the level in the tree for nodes, given an array of parents.

Parameters:

parents (np.ndarray) – Numpy array of integers with first entry -1.

Returns:

Level of each node in the free excluding the first entry which is the root (level = 0).

Return type:

list[int]

level_state_prob(levels, num_states, tr_mat, init_prob, out)[source]

Numba kernel: marginal state probabilities per tree level, out[k] = init_prob @ tr_mat^k.

class VariationalEmbeddingAttentionDistribution(mean, log_var, emission, position_prior, sigma2=0.5, name=None, optimizer_state=None)[source]

Bases: SequenceEncodableProbabilityDistribution

Responsibility-attention head over tied latent embeddings (mean-field posterior).

Parameters:

  • mean (np.ndarray)

  • log_var (np.ndarray)

  • emission (np.ndarray)

  • position_prior (np.ndarray)

  • sigma2 (float)

  • name (str | None)

  • optimizer_state (AttentionOptimizerState | dict[str, Any] | None)

density(x)[source]

Return the posterior-mean probability of one attention observation.

Parameters:

x (tuple[Any, int, int])

Return type:

float

log_density(x)[source]

Plug-in (posterior-mean) log conditional density log p(target | context, query).

Parameters:

x (tuple[Any, int, int])

Return type:

float

seq_log_density(x)[source]

Posterior-mean log density over an encoded batch -> (n,) (uses e = m).

Parameters:

x (tuple[ndarray, ndarray, ndarray])

Return type:

ndarray

predict_proba(context, query)[source]

Predictive target distribution from posterior-mean attention; (T,) or (n, T).

Parameters:
Return type:

ndarray

embeddings()[source]

The learned (posterior-mean) tied embedding table (S, D).

Return type:

ndarray

sampler(seed=None)[source]

Sample iid observations from the posterior-mean plug-in law used by scoring.

Parameters:

seed (int | None)

Return type:

VariationalEmbeddingAttentionSampler

posterior_predictive_sampler(seed=None)[source]

Sample iid observations, independently drawing an embedding table for every event.

Parameters:

seed (int | None)

Return type:

VariationalEmbeddingAttentionSampler

estimator(pseudo_count=None)[source]

Return a variational EM estimator initialized with this model’s dimensions.

Parameters:

pseudo_count (float | None)

Return type:

VariationalEmbeddingAttentionEstimator

dist_to_encoder()[source]

Return the encoder for contexts, queries, and targets.

Return type:

VariationalEmbeddingAttentionDataEncoder

class VariationalEmbeddingAttentionSampler(dist, seed=None, *, posterior_predictive=False)[source]

Bases: DistributionSampler

Sampler for either the plug-in or explicitly requested posterior-predictive law.

Parameters:

  • dist (VariationalEmbeddingAttentionDistribution)

  • seed (int | None)

  • posterior_predictive (bool)

sample(size=None, *, batched=True)[source]

Draw one observation or size iid synthetic observations.

Parameters:
Return type:

Any

class VariationalEmbeddingAttentionAccumulator(num_symbols, context_length, embed_dim, num_targets, mc, seed, keys=None, name=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulates the additive ELBO-gradient + emission/position sufficient statistics.

seq_update(x, weights, estimate)[source]

Update ELBO gradients and closed-form count statistics from encoded observations.

Parameters:

estimate (VariationalEmbeddingAttentionDistribution)

Return type:

None

update(x, weight, estimate)[source]

Update from one weighted latent-embedding attention observation.

Parameters:

weight (float)

Return type:

None

seq_initialize(x, weights, rng)[source]

Initialize emission and position counts with random attention responsibilities.

Parameters:

rng (RandomState)

Return type:

None

initialize(x, weight, rng)[source]

Initialize from one weighted latent-embedding attention observation.

Parameters:
Return type:

None

combine(suff_stat)[source]

Merge ELBO gradients, emission counts, position counts, and scalar totals.

Return type:

VariationalEmbeddingAttentionAccumulator

value()[source]

Return accumulated gradients, count statistics, log-likelihood, and total weight.

from_value(x)[source]

Restore accumulator state from value output.

Return type:

VariationalEmbeddingAttentionAccumulator

acc_to_encoder()[source]

Return the encoder compatible with this accumulator.

Return type:

VariationalEmbeddingAttentionDataEncoder

class VariationalEmbeddingAttentionAccumulatorFactory(estimator, keys=None, name=None)[source]

Bases: StatisticAccumulatorFactory

Create accumulators for variational embedding attention EM steps.

Parameters:

estimator (VariationalEmbeddingAttentionEstimator)

make()[source]

Create an accumulator with a deterministic per-iteration Monte-Carlo seed.

Return type:

VariationalEmbeddingAttentionAccumulator

class VariationalEmbeddingAttentionEstimator(num_symbols, context_length, embed_dim, num_targets, *, sigma2=0.5, lr=0.05, mc=6, prior_strength=1.0, emission_smoothing=1e-4, seed=0, name=None, keys=None)[source]

Bases: ParameterEstimator

Stateless variational-EM estimator consuming explicit model-carried Adam state.

Parameters:
accumulator_factory()[source]

Return a factory for variational embedding attention accumulators.

Return type:

VariationalEmbeddingAttentionAccumulatorFactory

estimate(nobs, suff_stat)[source]

Apply one variational EM update and return the updated attention distribution.

Parameters:

nobs (float | None)

Return type:

VariationalEmbeddingAttentionDistribution

class VariationalEmbeddingAttentionDataEncoder(num_symbols=None, context_length=None, num_targets=None)[source]

Bases: DataSequenceEncoder

Encodes (context_symbols, query_symbol, target) triples into stacked integer arrays.

seq_encode(x)[source]

Encode (context_symbols, query, target) observations.

Parameters:

x (Sequence[tuple[Any, int, int]])

Return type:

tuple[ndarray, ndarray, ndarray]

row_count(x)[source]

Return the number of observations represented by x.

The default covers simple array/list encodings and structurally aligned tuples/dicts. Encoders with flattened, ragged, or otherwise specialized layouts must override this method so metadata wrappers can verify counts instead of trusting a caller assertion.

Return type:

int

class VariationalMultiHopAttentionDistribution(mean, log_var, emission, sigma2=0.3, name=None, optimizer_state=None)[source]

Bases: SequenceEncodableProbabilityDistribution

A 2-hop chain over tied latent embeddings (mean-field posterior).

Parameters:

  • sigma2 (float)

  • name (str | None)

  • optimizer_state (AttentionOptimizerState | dict[str, Any] | None)

density(x)[source]

Return the probability of one context/query/target observation.

Return type:

float

log_density(x)[source]

Return the log-probability of one context/query/target observation.

Return type:

float

seq_log_density(x)[source]

Return vectorized log-probabilities for encoded two-hop attention observations.

Return type:

ndarray

predict_proba(context_keys, context_values, query)[source]

Predictive target distribution (posterior-mean embeddings); (T,) or (n, T).

Return type:

ndarray

embeddings()[source]

Return posterior mean embeddings for the tied latent symbols.

Return type:

ndarray

sampler(seed=None)[source]

Sample iid observations from the posterior-mean plug-in law used by scoring.

Parameters:

seed (int | None)

Return type:

VariationalMultiHopAttentionSampler

posterior_predictive_sampler(seed=None)[source]

Sample iid observations, independently drawing an embedding table for every event.

Parameters:

seed (int | None)

Return type:

VariationalMultiHopAttentionSampler

estimator(pseudo_count=None)[source]

Return a variational EM estimator initialized with this model’s dimensions.

Parameters:

pseudo_count (float | None)

Return type:

VariationalMultiHopAttentionEstimator

dist_to_encoder()[source]

Return the encoder for context keys, values, query symbols, and targets.

Return type:

VariationalMultiHopAttentionDataEncoder

class VariationalMultiHopAttentionSampler(dist, seed=None, *, posterior_predictive=False)[source]

Bases: DistributionSampler

Sampler for either the plug-in or explicitly requested posterior-predictive law.

Parameters:

  • seed (int | None)

  • posterior_predictive (bool)

sample(size=None, *, batched=True)[source]

Draw one observation or size iid synthetic observations.

Parameters:
Return type:

Any

class VariationalMultiHopAttentionAccumulator(num_symbols, embed_dim, num_targets, mc, seed, keys=None, name=None)[source]

Bases: SequenceEncodableStatisticAccumulator

Accumulate Monte-Carlo ELBO gradients and emission responsibilities for variational EM.

seq_update(x, weights, estimate)[source]

Update ELBO gradients and emission counts from encoded observations.

Return type:

None

seq_initialize(x, weights, rng)[source]

Initialize emission counts with random final-hop responsibilities.

Parameters:

rng (RandomState)

Return type:

None

update(x, weight, estimate)[source]

Update from one weighted two-hop attention observation.

Return type:

None

initialize(x, weight, rng)[source]

Initialize from one weighted two-hop attention observation.

Return type:

None

combine(suff_stat)[source]

Merge variational gradients, emission counts, log-likelihood, and weight totals.

value()[source]

Return accumulated gradients, emission counts, log-likelihood, and total weight.

from_value(x)[source]

Restore accumulator state from value output.

acc_to_encoder()[source]

Return the encoder compatible with this attention accumulator.

class VariationalMultiHopAttentionAccumulatorFactory(estimator, keys=None, name=None)[source]

Bases: StatisticAccumulatorFactory

Create accumulators for variational multi-hop attention EM steps.

make()[source]

Create an accumulator with a deterministic per-iteration Monte-Carlo seed.

class VariationalMultiHopAttentionEstimator(num_symbols, embed_dim, num_targets, *, sigma2=0.3, lr=0.05, mc=5, prior_strength=0.1, anneal_iters=100, emission_smoothing=1e-4, seed=0, name=None, keys=None)[source]

Bases: ParameterEstimator

Stateless variational-EM estimator with model-carried Adam state and prior annealing.

Parameters:
accumulator_factory()[source]

Return a factory for variational multi-hop attention accumulators.

estimate(nobs, suff_stat)[source]

Apply one variational EM update and return the updated attention distribution.

class VariationalMultiHopAttentionDataEncoder(num_symbols=None, num_targets=None)[source]

Bases: DataSequenceEncoder

Encode context keys, context values, query symbols, and targets as integer arrays.

seq_encode(x)[source]

Encode (context_keys, context_values, query, target) observations.

Parameters:

x (Sequence[tuple[Any, Any, int, int]])

row_count(x)[source]

Return the number of observations represented by x.

The default covers simple array/list encodings and structurally aligned tuples/dicts. Encoders with flattened, ragged, or otherwise specialized layouts must override this method so metadata wrappers can verify counts instead of trusting a caller assertion.

Return type:

int

Submodules