Repository: pyro-ppl/pyro

Issue:
FR parallel-scan GaussianHMM.rsample()
`GaussianHMM.rsample()` would enable a powerful class of guides for time-series models with non-Gaussian likelihoods. Consider for example a guide:
```py
from pyro.distributions import GaussianHMM
from pyro.ops.gaussian import Gaussian

def model(data):
    p_z = GaussianHMM(*my_prior)  # prior, say a Matern GP
    z = pyro.sample("z", p_z)  # samples from prior
    p_x = dist.Normal(*noisy_nonlinear_function(z))
    pyro.sample("x", p_x, obs=data)  # likelihood

def guide(data):
    p_z = GaussianHMM(*my_prior)  # same as in the model
    p_x_given_z = Gaussian(*my_nn(data))  # independent over time
    p_z_given_x = forward_filter(p_z, p_x_given_z)  # updates trans_dist
    assert isinstance(p_z_given_x, GaussianHMM)  # posterior is still Gaussian
    pyro.sample("z", p_z_given_x)  # samples from posterior via GaussianHMM.rsample()
```
In this guide, `my_nn` acts as a weak encoder, forming a gaussian update `(log_normalizer, info_vec, precision)` independently at each time step, kind of a neural EP. In principle it is easy to fuse these weak likelihoods with a strong prior, e.g. a Matern GP, and then draw a posterior sample using parallel-scan.

This guide would work well with either continuous observations and `Trace_ELBO` loss, or discrete observations (e.g. `Poisson`) and a `Trace_CRPS` loss. I'd like to add this class of guides to `pyro.contrib.timeseries`.


Discussion:

Resolving patch:
diff --git a/pyro/distributions/hmm.py b/pyro/distributions/hmm.py
--- a/pyro/distributions/hmm.py
+++ b/pyro/distributions/hmm.py
@@ -3,7 +3,6 @@
 
 import torch
 from torch.distributions import constraints
-from torch.distributions.utils import lazy_property
 
 from pyro.distributions.torch import Categorical, Gamma, Independent, MultivariateNormal
 from pyro.distributions.torch_distribution import TorchDistribution
@@ -90,6 +89,82 @@ def _sequential_gaussian_tensordot(gaussian):
     return gaussian[..., 0]
 
 
+def _is_subshape(x, y):
+    return broadcast_shape(x, y) == y
+
+
+def _sequential_gaussian_filter_sample(init, trans, sample_shape):
+    """
+    Draws a reparameterized sample from a Markov product of Gaussians via
+    parallel-scan forward-filter backward-sample.
+    """
+    assert isinstance(init, Gaussian)
+    assert isinstance(trans, Gaussian)
+    assert trans.dim() == 2 * init.dim()
+    assert _is_subshape(trans.batch_shape[:-1], init.batch_shape)
+    state_dim = trans.dim() // 2
+    device = trans.precision.device
+    perm = torch.cat([torch.arange(1 * state_dim, 2 * state_dim, device=device),
+                      torch.arange(0 * state_dim, 1 * state_dim, device=device),
+                      torch.arange(2 * state_dim, 3 * state_dim, device=device)])
+
+    # Forward filter, similar to _sequential_gaussian_tensordot().
+    tape = []
+    shape = trans.batch_shape[:-1]  # Note trans may be unbroadcasted.
+    gaussian = trans
+    while gaussian.batch_shape[-1] > 1:
+        time = gaussian.batch_shape[-1]
+        even_time = time // 2 * 2
+        even_part = gaussian[..., :even_time]
+        x_y = even_part.reshape(shape + (even_time // 2, 2))
+        x, y = x_y[..., 0], x_y[..., 1]
+        x = x.event_pad(right=state_dim)
+        y = y.event_pad(left=state_dim)
+        joint = (x + y).event_permute(perm)
+        tape.append(joint)
+        contracted = joint.marginalize(left=state_dim)
+        if time > even_time:
+            contracted = Gaussian.cat((contracted, gaussian[..., -1:]), dim=-1)
+        gaussian = contracted
+    gaussian = gaussian[..., 0] + init.event_pad(right=state_dim)
+
+    # Backward sample.
+    shape = sample_shape + init.batch_shape
+    result = gaussian.rsample(sample_shape).reshape(shape + (2, state_dim))
+    for joint in reversed(tape):
+        # The following comments demonstrate two example computations, one
+        # EVEN, one ODD.  Ignoring sample_shape and batch_shape, let each zn be
+        # a single sampled event of shape (state_dim,).
+        if joint.batch_shape[-1] == result.size(-2) - 1:  # EVEN case.
+            # Suppose e.g. result = [z0, z2, z4]
+            cond = result.repeat_interleave(2, dim=-2)  # [z0, z0, z2, z2, z4, z4]
+            cond = cond[..., 1:-1, :]  # [z0, z2, z2, z4]
+            cond = cond.reshape(shape + (-1, 2 * state_dim))  # [z0z2, z2z4]
+            sample = joint.condition(cond).rsample()  # [z1, z3]
+            sample = torch.nn.functional.pad(sample, (0, 0, 0, 1))  # [z1, z3, 0]
+            result = torch.stack([
+                result,  # [z0, z2, z4]
+                sample,  # [z1, z3, 0]
+            ], dim=-2)  # [[z0, z1], [z2, z3], [z4, 0]]
+            result = result.reshape(shape + (-1, state_dim))  # [z0, z1, z2, z3, z4, 0]
+            result = result[..., :-1, :]  # [z0, z1, z2, z3, z4]
+        else:  # ODD case.
+            assert joint.batch_shape[-1] == result.size(-2) - 2
+            # Suppose e.g. result = [z0, z2, z3]
+            cond = result[..., :-1, :].repeat_interleave(2, dim=-2)  # [z0, z0, z2, z2]
+            cond = cond[..., 1:-1, :]  # [z0, z2]
+            cond = cond.reshape(shape + (-1, 2 * state_dim))  # [z0z2]
+            sample = joint.condition(cond).rsample()  # [z1]
+            sample = torch.cat([sample, result[..., -1:, :]], dim=-2)  # [z1, z3]
+            result = torch.stack([
+                result[..., :-1, :],  # [z0, z2]
+                sample,  # [z1, z3]
+            ], dim=-2)  # [[z0, z1], [z2, z3]]
+            result = result.reshape(shape + (-1, state_dim))  # [z0, z1, z2, z3]
+
+    return result[..., 1:, :]  # [z1, z2, z3, ...]
+
+
 def _sequential_gamma_gaussian_tensordot(gamma_gaussian):
     """
     Integrates a GammaGaussian ``x`` whose rightmost batch dimension is time, computes::
@@ -318,39 +393,23 @@ def __init__(self, initial_dist, transition_matrix, transition_dist,
 
         self.hidden_dim = hidden_dim
         self.obs_dim = obs_dim
-        self.initial_dist = initial_dist
-        self.transition_matrix = transition_matrix
-        self.transition_dist = transition_dist
-        self.observation_matrix = observation_matrix
-        self.observation_dist = observation_dist
-
-    @lazy_property
-    def _init(self):
-        # To save computation in _sequential_gaussian_tensordot(), we expand
-        # only _init, which is applied only after
-        # _sequential_gaussian_tensordot().
-        return mvn_to_gaussian(self.initial_dist).expand(self.batch_shape)
-
-    @lazy_property
-    def _trans(self):
-        return matrix_and_mvn_to_gaussian(self.transition_matrix, self.transition_dist)
-
-    @lazy_property
-    def _obs(self):
-        return matrix_and_mvn_to_gaussian(self.observation_matrix, self.observation_dist)
+        self._init = mvn_to_gaussian(initial_dist).expand(batch_shape)
+        self._trans = matrix_and_mvn_to_gaussian(transition_matrix, transition_dist)
+        self._obs = matrix_and_mvn_to_gaussian(observation_matrix, observation_dist)
 
     def expand(self, batch_shape, _instance=None):
         new = self._get_checked_instance(GaussianHMM, _instance)
         new.hidden_dim = self.hidden_dim
         new.obs_dim = self.obs_dim
-        new.initial_dist = self.initial_dist
-        new.transition_matrix = self.transition_matrix
-        new.transition_dist = self.transition_dist
-        new.observation_matrix = self.observation_matrix
-        new.observation_dist = self.observation_dist
+        new._obs = self._obs
+        new._trans = self._trans
 
-        # Expand lazily.
+        # To save computation in _sequential_gaussian_tensordot(), we expand
+        # only _init, which is applied only after
+        # _sequential_gaussian_tensordot().
         batch_shape = torch.Size(broadcast_shape(self.batch_shape, batch_shape))
+        new._init = self._init.expand(batch_shape)
+
         super(GaussianHMM, new).__init__(batch_shape, self.event_shape, validate_args=False)
         new._validate_args = self.__dict__.get('_validate_args')
         return new
@@ -370,14 +429,15 @@ def log_prob(self, value):
         return result
 
     def rsample(self, sample_shape=torch.Size()):
-        batch_shape = self.batch_shape
-        time_shape = self.event_shape[:1]
-        init = self.initial_dist.expand(batch_shape).rsample(sample_shape)
-        trans = self.transition_dist.expand(batch_shape + time_shape).rsample(sample_shape)
-        obs = self.observation_dist.expand(batch_shape + time_shape).rsample(sample_shape)
-        mat = self.transition_matrix.expand(batch_shape + time_shape + (self.hidden_dim, self.hidden_dim))
-        z = _linear_integrate(init, mat, trans)
-        return (z.unsqueeze(-2) @ self.observation_matrix).squeeze(-2) + obs
+        sample_shape = torch.Size(sample_shape)
+        obs_dim = self.obs_dim
+        hidden_dim = self.hidden_dim
+        obs = self._obs.marginalize(right=self.obs_dim).event_pad(left=self.hidden_dim)
+        z = _sequential_gaussian_filter_sample(self._init, self._trans + obs, sample_shape)
+        perm = torch.cat([torch.arange(hidden_dim, hidden_dim + obs_dim, device=z.device),
+                          torch.arange(hidden_dim, device=z.device)])
+        x = self._obs.event_permute(perm).condition(z).rsample()
+        return x
 
     def filter(self, value):
         """
@@ -405,6 +465,48 @@ def filter(self, value):
         return MultivariateNormal(loc, precision_matrix=precision,
                                   validate_args=self._validate_args)
 
+    def conjugate_update(self, other):
+        """
+        Creates an updated :class:`GaussianHMM` fusing information from another
+        compatible distribution.
+
+        This should satisfy::
+
+            fg, log_normalizer = f.conjugate_update(g)
+            assert f.log_prob(x) + g.log_prob(x) == fg.log_prob(x) + log_normalizer
+
+        :param other: A distribution representing ``p(data|self.probs)`` but
+            normalized over ``self.probs`` rather than ``data``.
+        :type other: ~torch.distributions.Independent of
+            ~torch.distributions.MultivariateNormal or ~torch.distributions.Normal
+        :return: a pair ``(updated,log_normalizer)`` where ``updated`` is an
+            updated :class:`GaussianHMM` , and ``log_normalizer`` is a
+            :class:`~torch.Tensor` representing the normalization factor.
+        """
+        assert other.event_shape == self.event_shape
+        assert (isinstance(other, torch.distributions.Independent) and
+                (isinstance(other.base_dist, torch.distributions.Normal) or
+                 isinstance(other.base_dist, torch.distributions.MultivariateNormal)))
+        new = self._get_checked_instance(GaussianHMM)
+        new.hidden_dim = self.hidden_dim
+        new.obs_dim = self.obs_dim
+        new._init = self._init
+        new._trans = self._trans
+        new._obs = self._obs + mvn_to_gaussian(other.to_event(-1)).event_pad(left=self.hidden_dim)
+
+        # Normalize.
+        # TODO cache this computation for the forward pass of .rsample().
+        logp = new._trans + new._obs.marginalize(right=new.obs_dim).event_pad(left=new.hidden_dim)
+        logp = _sequential_gaussian_tensordot(logp.expand(logp.batch_shape))
+        logp = gaussian_tensordot(new._init, logp, dims=new.hidden_dim)
+        log_normalizer = logp.event_logsumexp()
+        new._init = new._init - log_normalizer
+
+        batch_shape = log_normalizer.shape
+        super(GaussianHMM, new).__init__(batch_shape, self.event_shape, validate_args=False)
+        new._validate_args = self.__dict__.get('_validate_args')
+        return new, log_normalizer
+
 
 class GammaGaussianHMM(TorchDistribution):
     """
diff --git a/pyro/ops/gaussian.py b/pyro/ops/gaussian.py
--- a/pyro/ops/gaussian.py
+++ b/pyro/ops/gaussian.py
@@ -103,11 +103,19 @@ def __add__(self, other):
         """
         Adds two Gaussians in log-density space.
         """
-        assert isinstance(other, Gaussian)
-        assert self.dim() == other.dim()
-        return Gaussian(self.log_normalizer + other.log_normalizer,
-                        self.info_vec + other.info_vec,
-                        self.precision + other.precision)
+        if isinstance(other, Gaussian):
+            assert self.dim() == other.dim()
+            return Gaussian(self.log_normalizer + other.log_normalizer,
+                            self.info_vec + other.info_vec,
+                            self.precision + other.precision)
+        if isinstance(other, (int, float, torch.Tensor)):
+            return Gaussian(self.log_normalizer + other, self.info_vec, self.precision)
+        raise ValueError("Unsupported type: {}".format(type(other)))
+
+    def __sub__(self, other):
+        if isinstance(other, (int, float, torch.Tensor)):
+            return Gaussian(self.log_normalizer - other, self.info_vec, self.precision)
+        raise ValueError("Unsupported type: {}".format(type(other)))
 
     def log_density(self, value):
         """
@@ -242,7 +250,6 @@ class AffineNormal:
     """
     def __init__(self, matrix, loc, scale):
         assert loc.shape == scale.shape
-        x_dim, y_dim = matrix.shape[-2:]
         self.matrix = matrix
         self.loc = loc
         self.scale = scale
@@ -267,9 +274,15 @@ def to_gaussian(self):
         mvn = torch.distributions.MultivariateNormal(self.loc, scale_tril=self.scale.diag_embed())
         return matrix_and_mvn_to_gaussian(self.matrix, mvn)
 
+    def event_permute(self, perm):
+        return self.to_gaussian().event_permute(perm)
+
     def __add__(self, other):
         return self.to_gaussian() + other
 
+    def marginalize(self, left=0, right=0):
+        return self.to_gaussian().marginalize(left, right)
+
 
 def mvn_to_gaussian(mvn):
     """


Test patch:
diff --git a/tests/distributions/test_hmm.py b/tests/distributions/test_hmm.py
--- a/tests/distributions/test_hmm.py
+++ b/tests/distributions/test_hmm.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2017-2019 Uber Technologies, Inc.
 # SPDX-License-Identifier: Apache-2.0
 
+import logging
 import operator
 from functools import reduce
 
@@ -10,17 +11,19 @@
 
 import pyro
 import pyro.distributions as dist
-from pyro.distributions.hmm import (_sequential_gamma_gaussian_tensordot,
+from pyro.distributions.hmm import (_sequential_gamma_gaussian_tensordot, _sequential_gaussian_filter_sample,
                                     _sequential_gaussian_tensordot, _sequential_logmatmulexp)
 from pyro.distributions.util import broadcast_shape
 from pyro.infer import TraceEnum_ELBO, config_enumerate
+from pyro.ops.gamma_gaussian import (gamma_and_mvn_to_gamma_gaussian, gamma_gaussian_tensordot,
+                                     matrix_and_mvn_to_gamma_gaussian)
 from pyro.ops.gaussian import gaussian_tensordot, matrix_and_mvn_to_gaussian, mvn_to_gaussian
-from pyro.ops.gamma_gaussian import (gamma_gaussian_tensordot, matrix_and_mvn_to_gamma_gaussian,
-                                     gamma_and_mvn_to_gamma_gaussian)
 from pyro.ops.indexing import Vindex
 from tests.common import assert_close
-from tests.ops.gaussian import assert_close_gaussian, random_gaussian, random_mvn
 from tests.ops.gamma_gaussian import assert_close_gamma_gaussian, random_gamma, random_gamma_gaussian
+from tests.ops.gaussian import assert_close_gaussian, random_gaussian, random_mvn
+
+logger = logging.getLogger(__name__)
 
 
 def check_expand(old_dist, old_data):
@@ -73,6 +76,17 @@ def test_sequential_gaussian_tensordot(batch_shape, state_dim, num_steps):
     assert_close_gaussian(actual, expected)
 
 
+@pytest.mark.parametrize('num_steps', list(range(1, 20)))
+@pytest.mark.parametrize('state_dim', [1, 2, 3])
+@pytest.mark.parametrize('batch_shape', [(), (4,), (3, 2)], ids=str)
+@pytest.mark.parametrize('sample_shape', [(), (4,), (3, 2)], ids=str)
+def test_sequential_gaussian_filter_sample(sample_shape, batch_shape, state_dim, num_steps):
+    init = random_gaussian(batch_shape, state_dim)
+    trans = random_gaussian(batch_shape + (num_steps,), state_dim + state_dim)
+    sample = _sequential_gaussian_filter_sample(init, trans, sample_shape)
+    assert sample.shape == sample_shape + batch_shape + (num_steps, state_dim)
+
+
 @pytest.mark.parametrize('num_steps', list(range(1, 20)))
 @pytest.mark.parametrize('state_dim', [1, 2, 3])
 @pytest.mark.parametrize('batch_shape', [(), (5,), (2, 4)], ids=str)
@@ -289,6 +303,17 @@ def test_gaussian_hmm_shape(diag, init_shape, trans_mat_shape, trans_mvn_shape,
     x = d.expand((6, 5)).rsample()
     assert x.shape == (6, 5) + d.event_shape
 
+    likelihood = dist.Normal(data, 1).to_event(2)
+    p, log_normalizer = d.conjugate_update(likelihood)
+    assert p.batch_shape == d.batch_shape
+    assert p.event_shape == d.event_shape
+    x = p.rsample()
+    assert x.shape == d.shape()
+    x = p.rsample((6,))
+    assert x.shape == (6,) + d.shape()
+    x = p.expand((6, 5)).rsample()
+    assert x.shape == (6, 5) + d.event_shape
+
     final = d.filter(data)
     assert isinstance(final, dist.MultivariateNormal)
     assert final.batch_shape == d.batch_shape
@@ -327,11 +352,14 @@ def test_gaussian_hmm_distribution(diag, sample_shape, batch_shape, num_steps, h
     #    init | H
     #   trans | H H H H            H = hidden
     #     obs |   H H H O O O      O = observed
+    #    like |         O O O
     # and then combine these using gaussian_tensordot().
     T = num_steps
     init = mvn_to_gaussian(init_dist)
     trans = matrix_and_mvn_to_gaussian(trans_mat, trans_dist)
     obs = matrix_and_mvn_to_gaussian(obs_mat, obs_mvn)
+    like_dist = dist.Normal(torch.randn(data.shape), 1).to_event(2)
+    like = mvn_to_gaussian(like_dist)
 
     unrolled_trans = reduce(operator.add, [
         trans[..., t].event_pad(left=t * hidden_dim, right=(T - t - 1) * hidden_dim)
@@ -341,6 +369,10 @@ def test_gaussian_hmm_distribution(diag, sample_shape, batch_shape, num_steps, h
         obs[..., t].event_pad(left=t * obs.dim(), right=(T - t - 1) * obs.dim())
         for t in range(T)
     ])
+    unrolled_like = reduce(operator.add, [
+        like[..., t].event_pad(left=t * obs_dim, right=(T - t - 1) * obs_dim)
+        for t in range(T)
+    ])
     # Permute obs from HOHOHO to HHHOOO.
     perm = torch.cat([torch.arange(hidden_dim) + t * obs.dim() for t in range(T)] +
                      [torch.arange(obs_dim) + hidden_dim + t * obs.dim() for t in range(T)])
@@ -355,26 +387,35 @@ def test_gaussian_hmm_distribution(diag, sample_shape, batch_shape, num_steps, h
     expected_log_prob = logp.log_density(unrolled_data)
     assert_close(actual_log_prob, expected_log_prob)
 
-    # Test mean and covariance.
+    d_posterior, log_normalizer = d.conjugate_update(like_dist)
+    assert_close(d.log_prob(data) + like_dist.log_prob(data),
+                 d_posterior.log_prob(data) + log_normalizer)
+
     if batch_shape or sample_shape:
         return
-    with torch.no_grad():
-        num_samples = 100000
-        samples = d.sample([num_samples]).reshape(num_samples, T * obs_dim)
-        actual_mean = samples.mean(0)
-        delta = samples - actual_mean
-        actual_cov = (delta.unsqueeze(-1) * delta.unsqueeze(-2)).mean(0)
-        actual_scale = actual_cov.diagonal(dim1=-2, dim2=-1).sqrt()
-        actual_corr = actual_cov / (actual_scale.unsqueeze(-1) * actual_scale.unsqueeze(-2))
-
-        expected_cov = logp.precision.cholesky().cholesky_inverse()
-        expected_mean = expected_cov.matmul(logp.info_vec.unsqueeze(-1)).squeeze(-1)
-        expected_scale = expected_cov.diagonal(dim1=-2, dim2=-1).sqrt()
-        expected_corr = expected_cov / (expected_scale.unsqueeze(-1) * expected_scale.unsqueeze(-2))
-
-        assert_close(actual_mean, expected_mean, atol=0.02, rtol=0.02)
-        assert_close(actual_scale, expected_scale, atol=0.02, rtol=0.02)
-        assert_close(actual_corr, expected_corr, atol=0.01)
+
+    # Test mean and covariance.
+    prior = "prior", d, logp
+    posterior = "posterior", d_posterior, logp + unrolled_like
+    for name, d, g in [prior, posterior]:
+        logging.info("testing {} moments".format(name))
+        with torch.no_grad():
+            num_samples = 100000
+            samples = d.sample([num_samples]).reshape(num_samples, T * obs_dim)
+            actual_mean = samples.mean(0)
+            delta = samples - actual_mean
+            actual_cov = (delta.unsqueeze(-1) * delta.unsqueeze(-2)).mean(0)
+            actual_std = actual_cov.diagonal(dim1=-2, dim2=-1).sqrt()
+            actual_corr = actual_cov / (actual_std.unsqueeze(-1) * actual_std.unsqueeze(-2))
+
+            expected_cov = g.precision.cholesky().cholesky_inverse()
+            expected_mean = expected_cov.matmul(g.info_vec.unsqueeze(-1)).squeeze(-1)
+            expected_std = expected_cov.diagonal(dim1=-2, dim2=-1).sqrt()
+            expected_corr = expected_cov / (expected_std.unsqueeze(-1) * expected_std.unsqueeze(-2))
+
+            assert_close(actual_mean, expected_mean, atol=0.05, rtol=0.02)
+            assert_close(actual_std, expected_std, atol=0.05, rtol=0.02)
+            assert_close(actual_corr, expected_corr, atol=0.02)
 
 
 @pytest.mark.parametrize('obs_dim', [1, 2, 3])