.

torch_struct/distributions.py

@@ -194,6 +194,26 @@ class LinearChainCRF(StructDistribution):
 
     struct = LinearChain
 
+class HMM(StructDistribution):
+    r"""
+    Represents hidden-markov smoothing with C hidden states.
+
+    Event shape is of the form:
+
+    Parameters:
+        transition: C X C
+        emission: V x C
+        init: C
+        observations: b x N between [0, V-1]
+
+    Compact representation: N long tensor in [0, ..., C-1]
+    """
+    def __init__(self, transitions, emission, init, observations, lengths=None):
+        log_potentials = HMM.struct.hmm(transition, emission, init, observations)
+        super().__init__(log_potentials, lengths)
+
+    struct = LinearChain
+
 
 class SemiMarkovCRF(StructDistribution):
     r"""

torch_struct/linearchain.py

@@ -1,3 +1,22 @@
+"""
+
+A linear-chain dynamic program.
+
+Considers parameterized functions of the form :math:`f: {\cal Y} \rightarrow \mathbb{R}`.
+
+Combinatorial set :math:`{y_{1:N} \in \cal Y}` with each :math:`y_n \in {1, \ldots, C}`
+
+Function factors as :math:`f(y) = \prod_{n=1}^N \phi(n, y_n, y_n{-1})`
+
+Example use cases:
+
+* Part-of-Speech Tagging
+* Sequence Labeling
+* Hidden Markov Models
+
+"""
+
+
 import torch
 from .helpers import _Struct
 import math