mixle.stats.latent.hidden_markov module

Hidden Markov models with state-specific emission distributions.

Data type: Sequence[T] (determined by emission distributions).

Consider an observation x = (x_1, x_2, …, x_T) where x_i is of data type T. Assume Z = (Z_1, …, Z_T) is an unobserved sequence of hidden states taking on values {1,2,..,K}. A K state hidden markov model can be written as hierarchical model as follows:

For t = 1,2,..,T, the emission distributions are given by

1. P_1(X_t = x_t | Z_t = k), for k = {1,2,…,K}.

The state transitions are given by the K by K matrix formed from

2. p_mat(Z_t = i | Z_{t-1} = j), for i, j = {2,3,..,K}.

The initial state distribution is given by weights

3. p_mat(Z_1=k) = pi_k, for k = {1,2,…,K}, where sum_k pi_k = 1.0

If included, the length of the hidden markov model sequences is modeled through

4. P_len(T), where P_len() is a distribution with support on non-negative integers.

The emission distributions must accept observations of type T. The transition matrix is row-stochastic, and the initial weight vector sums to one.

hmm_dirichlet_default_prior(num_states)

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)

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:

• log_w (ndarray)

• log_a (ndarray)

• log_b (ndarray)

• term_mask (ndarray)

Return type:

float

terminal_forward_backward(log_w, log_a, log_b, term_mask)

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:

• log_w (ndarray)

• log_a (ndarray)

• log_b (ndarray)

• term_mask (ndarray)

Return type:

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

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

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)

get_prior()

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)

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)

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)

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

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

compute_declaration()

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

density(x)

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)

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)|x_mat(0),…,x_mat(t-1)) = p_mat(x_mat(t)|Z(t))*p_mat(Z(t)|Z(t-1))*p_mat(Z(t-1)|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)

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)

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:

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

• engine (Any)

Return type:

Any

seq_posterior(x)

Return vectorized posterior state probabilities for encoded observations.

Parameters:

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

Return type:

list[ndarray] | None

viterbi(x)

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

Parameters:

x (list[T])

Return type:

ndarray

latent_posterior(x)

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

Parameters:

x (list[T])

Return type:

MarkovChainLatentPosterior

posterior_predictive(x, seed=None)

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:

• x (list[T])

• seed (int | None)

Return type:

list[Any]

seq_viterbi(x)

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

The result is indexed in the encoder’s flattened sequence layout. Use the data encoder metadata to map flattened assignments back to individual sequences when working directly with encoded data.

Parameters:

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

to_fisher(**kwargs)

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

density_semantics()

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

sampler(seed=None)

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)

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

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

Return an encoder for HMM observation sequences and optional lengths.

Return type:

HiddenMarkovDataEncoder

enumerator()

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

Return type:

HiddenMarkovModelEnumerator

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

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)

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)

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)

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)

Bases: DistributionSampler

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

Parameters:

• dist (HiddenMarkovModelDistribution)

• seed (int | None)

sample_seq(size=None, *, batched=True)

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)

Sample one HMM sequence until an emitted terminal value appears.

Parameters:

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

Returns:

A sequence whose final value belongs to terminal_set.

Return type:

list[T]

sample_terminal_states(cap=1_000_000)

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:

cap (int)

Return type:

list[T]

sample(size=None, *, batched=True)

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=False, keys=(None, None, None), name=None)

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)

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)

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)

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)

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)

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)

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

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)

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)

Scale linear HMM sufficient statistics while preserving metadata.

Parameters:

c (float)

Return type:

HiddenMarkovAccumulator

key_merge(stats_dict)

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)

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

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

Return type:

HiddenMarkovDataEncoder

class HiddenMarkovAccumulatorFactory(factories, len_factory=NullAccumulatorFactory(), use_numba=False, keys=(None, None, None), name=None)

Bases: StatisticAccumulatorFactory

Factory for HMM Baum-Welch sufficient-statistic accumulators.

Parameters:

• factories (Sequence[StatisticAccumulatorFactory])

• len_factory (StatisticAccumulatorFactory)

• use_numba (bool)

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

• name (str | None)

make()

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)

Bases: ParameterEstimator

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

Parameters:

• estimators (list[ParameterEstimator])

• len_estimator (ParameterEstimator | None)

• pseudo_count (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()

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

get_prior()

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)

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)

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)

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:

• nobs (Optional[float]) – Number of observations used in estimation.

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

Returns:

HiddenMarkovModelDistribution object.

Return type:

HiddenMarkovModelDistribution

class HiddenMarkovDataEncoder(emission_encoder, len_encoder=NullDataEncoder(), use_numba=False)

Bases: DataSequenceEncoder

Encoder for HMM observation sequences and optional sequence lengths.

Parameters:

• emission_encoder (DataSequenceEncoder)

• len_encoder (DataSequenceEncoder | None)

• use_numba (bool)

seq_encode(x)

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]

vec_bincount1(x, w, out)

Numba bincount on the rows of matrix w for groups x.

Parameters:

• x (np.ndarray[np.float64]) – Group ids of rows

• w (np.ndarray[np.float64]) – N by S numpy array with rows corresponding to x

• out (np.ndarray[np.float64]) – Unique values in support of x by S.

Returns:

Numpy 2-d array.

vec_bincount2(x, w, out)

Numba bincount on the rows of matrix w for groups x.

N = len(x) S = number of states. U = unique values in x can take on.

Parameters:

• x (np.ndarray[np.float64]) – Group ids of columns of w.

• w (np.ndarray[np.float64]) – S by N numpy array with cols corresponding to x

• out (np.ndarray[np.float64]) – S by U matrix.

Returns:

Numpy 2-d array.

hmm_pad_log_emissions(log_emit_flat, sz)

Pack per-sequence-contiguous (tot, S) log-emissions into padded (N, Tmax, S) + (N, Tmax) mask.

Parameters:

• log_emit_flat (np.ndarray) – (tot, S) log emission densities, ordered sequence-by-sequence.

• sz (np.ndarray) – (N,) sequence lengths summing to tot.

Returns:

Tuple of (padded (N, Tmax, S) float64, mask (N, Tmax) float64, offsets (N+1,) int).

hmm_engine_emissions(topics, enc_data, engine, homogeneous=True)

Per-state log emissions scored ON the engine — an (tot, S) engine tensor — or None when any topic lacks engine scoring (the caller falls back to the host loop).

hmm_engine_pad_log_emissions(pr_dev, sz, engine)

Pack engine-resident (tot, S) emissions into padded (N, Tmax, S) + host (N, Tmax) mask.

The gather index is built host-side (low-overhead ints); the gather itself runs on the engine, so the emission matrix never round-trips through numpy on its way into the forward-backward.

hmm_engine_forward_ll(engine, log_emit, log_w, log_a, mask)

Per-sequence emission log-likelihood over padded sequences — the forward half of hmm_engine_forward_backward(), for scoring paths that don’t need posteriors.

Args mirror the forward-backward: log_emit (N, Tmax, S) padded log emissions, log_w (S,), log_a (S, S), mask (N, Tmax) 1.0 for real steps. Returns ll (N,).

hmm_engine_forward_backward(engine, log_emit, log_w, log_a, mask, weights=None)

Log-space forward-backward over padded sequences using ComputeEngine ops.

Parameters:

• engine (ComputeEngine) – Array backend (numpy or torch).

• log_emit – (N, Tmax, S) log emission densities; padded slots may be -inf (masked out).

• log_w – (S,) log initial-state probabilities.

• log_a – (S, S) log transition matrix.

• mask – (N, Tmax) 1.0 for real observations, 0.0 for padding.

• weights – Optional (N,) per-sequence weights applied to gamma/xi/pi (E-step reweighting).

Returns:

ll: (N,) per-sequence emission log-likelihood (add the length model separately). gamma: (N, Tmax, S) posterior state probabilities (weighted; padded slots zero). xi_sum: (S, S) expected transition counts over all sequences and steps (weighted). pi: (N, S) initial-state posteriors (weighted).

Return type:

Tuple of

HiddenMarkovModelAccumulator

alias of HiddenMarkovAccumulator

HiddenMarkovModelAccumulatorFactory

alias of HiddenMarkovAccumulatorFactory

HiddenMarkovModelDataEncoder

alias of HiddenMarkovDataEncoder

HiddenMarkovModelEstimator

alias of HiddenMarkovEstimator

HiddenMarkovModelSampler

alias of HiddenMarkovSampler

class HiddenMarkovFisherView(dist)

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)

Return Baum-Welch sufficient statistics for one weighted sequence.

Parameters:

• x (Any)

• estimate (Any | None)

• weight (float)

Return type:

Any

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

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

Parameters:

• stats (ndarray | None)

• diagonal (bool)

• ridge (float)

• kwargs (Any)

Return type:

ndarray

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

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:

• stats (ndarray | None)

• metric (str)

• center (ndarray | None)

• fisher (ndarray | None)

• ridge (float)

• kwargs (Any)

Return type:

ndarray