mixle.stats.latent.hidden_markov module

“Create, estimate, and sample from a hidden markov model with K emission distributions (i.e. K states).

Defines the HierarchicalMixtureDistribution, HierarchicalMixtureSampler, HierarchicalMixtureEstimatorAccumulatorFactory, HierarchicalMixtureEstimatorAccumulator, HierarchicalMixtureEstimator, and the HierarchicalMixtureDataEncoder classes for use with mixle.

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.

Note that P_1() in (1) must be a distribution compatible with type T data. p_mat() in (2) is a 2-d numpy array of 2-d list of floats where the rows sum to 1.0. (3) is represented by a numpy array of list of floats that sum to 1.

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

compute_declaration()

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 vectorized log-density values for sequence-encoded observations.

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 Viterbi paths for sequence-encoded observation sequences.

Parameters:

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

to_fisher(**kwargs)

Forward-backward Fisher view for the HMM.

density_semantics()

What log_density returns relative to the true log-density (default: exact).

Override to declare that this distribution’s log_density is a variational lower bound (ELBO), an upper bound, or an approximation rather than the exact log p(x). This is surfaced as the ExactDensity capability and noted in mixle.describe(), so code that needs an exact likelihood can require(x, ExactDensity) instead of silently trusting a bound.

sampler(seed=None)

Create a HiddenMarkovSampler object with seed passed.

Note: Throws exception if ‘len_dist’and ‘terminal_values’ are not set.

If len_dist is set, it should be a SequenceEncodableProbabilityDistribution with data type int and support on non-negative integers.

Parameters:

seed (Optional[int]) – Set seed for random sampling.

Returns:

HiddenMarkovSampler object.

Return type:

HiddenMarkovSampler

estimator(pseudo_count=None)

Create HiddenMarkovEstimator for estimating HiddenMarkovDistribution objects from aggregated sufficient

statistics.

Parameters:

pseudo_count (Optional[float]) – Used to re-weight sufficient statistics of HiddenMarkovDistribution object instance.

Returns:

HiddenMarkovEstimator object.

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

Returns HiddenMarkovDataEncoder object for encoding sequences of iid HMM observations.

Return type:

HiddenMarkovDataEncoder

enumerator()

Returns HiddenMarkovModelEnumerator iterating 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

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

Parameters:

• dist (HiddenMarkovModelDistribution)

• seed (int | None)

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

Sample iid HMM sequences.

If size is None, 1 sample is drawn and a List[T] is returned. If size > 0, ‘size’ samples are drawn and a List of length ‘size’ with HMM sequences (List[T]) is returned.

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 (Optional[int]) – 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:

List[T] or List[List[T]] depending on size arg.

Return type:

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

sample_terminal(terminal_set)

Sample an HMM sequence, until a terminal value is samples from the emission distribution.

Parameters:

terminal_set (Set[T]) – Set values to terminate the HMM sequence.

Returns:

List[T] with length determined by samples to reach the first terminating value.

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 samples from HMM.

If a ‘len_sampler’ is set, call ‘sample_seq()’ (See HiddenMarkovSampler.sample_seq() for details). If ‘len_sampler’ is the NullDistributionSampler(), ‘sample_terminal()’ is called. (See HiddenMarkovSampler.sample_terminal() for details).

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 (Optional[int]) – 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:

List[T] or List[List[T]] depending on arg size.

class HiddenMarkovAccumulator(accumulators, len_accumulator=NullAccumulator(), use_numba=False, keys=(None, None, None), name=None)

Bases: SequenceEncodableStatisticAccumulator

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: Note efficient. Should use seq_encode() for fully encoded sequence instead.

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 HiddenMarkovAccumulator object with HMM sequence x.

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 update for HiddenMarkovAccumulator object from encoded sequence of observations.

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

Returns sufficient statistics of HiddenMarkovAccumulator object instance.

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)

Set the sufficient statistics of HiddenMarkovAccumulator object instance to value x.

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

HiddenMarkovAccumulator object.

Return type:

HiddenMarkovAccumulator

scale(c)

Scale linear HMM sufficient statistics while preserving metadata.

Parameters:

c (float)

Return type:

HiddenMarkovAccumulator

key_merge(stats_dict)

Merge the sufficient statistics of object instance with sufficient statistics in suff_stat that have

matching keys.

Parameters:

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

Returns:

None.

Return type:

None

key_replace(stats_dict)

Replace the sufficient statistics of HiddenMarkovAccumulator object with matching sufficient statistics in

arg suff_stat that have matching keys.

Parameters:

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

Returns:

None.

Return type:

None

acc_to_encoder()

Returns HiddenMarkovDataEncoder object for encoding sequences of iid HMM observations.

Return type:

HiddenMarkovDataEncoder

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

Bases: StatisticAccumulatorFactory

Parameters:

• factories (Sequence[StatisticAccumulatorFactory])

• len_factory (StatisticAccumulatorFactory)

• use_numba (bool)

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

• name (str | None)

make()

Returns a HiddenMarkovAccumulator object.

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

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

Returns an HiddenMarkovAccumulatorFactory object.

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

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)

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

Structured sufficient stats for one observation.

For latent-variable distributions, this is the posterior-expected complete-data sufficient statistic under estimate (or this view’s distribution when estimate is omitted).

Parameters:

• x (Any)

• estimate (Any | None)

• weight (float)

Return type:

Any

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

Empirical Fisher approximation from per-observation statistic vectors.

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 centered/whitened sufficient-statistic vectors.

metric=’identity’ returns centered statistics, ‘diagonal’ divides by per-coordinate Fisher standard deviations, and ‘full’ applies an empirical full-matrix whitening transform.

Parameters:

• stats (ndarray | None)

• metric (str)

• center (ndarray | None)

• fisher (ndarray | None)

• ridge (float)

• kwargs (Any)

Return type:

ndarray