Repository: pyro-ppl/pyro

Issue:
FR StudentTReparam, StudentTHMM
This proposes to implement reparameterizers for `StudentT` distributions by introducing latent `Gamma`-distributed variables, following the pattern of of symmetric `Stable` distributions.

## Tasks
- [ ] `StudentTReparam` following [SymmetricStableReparam](http://docs.pyro.ai/en/latest/infer.reparam.html#pyro.infer.reparam.stable.SymmetricStableReparam). #2265 
- ~~`StudentTHMM` following [StableHMM](http://docs.pyro.ai/en/latest/distributions.html#pyro.distributions.StableHMM). Note this version should use batched `StudentT` distributions rather than the elliptical `MultivariateStudentT` distribution; the latter distribution might be separately generalized to an HMM.~~
    UPDATE this is obsoleted by the new generic `LinearHMM` #2260
- ~~`StudentTHMMReparam` following [StableHMMReparam](http://docs.pyro.ai/en/latest/infer.reparam.html#pyro.infer.reparam.stable.StableHMMReparam).~~
  UPDATE this is obsoleted by the new generic `LinearHMMReparam` class in #2259 


Discussion:
@fehiepsi I should have some time to implement this maybe next week, or did you want to implement it? It's nearly identical to the `Stable` versions.
@fritzo I can implement them later this week. Thanks for bringing them up!

Resolving patch:
diff --git a/pyro/infer/reparam/__init__.py b/pyro/infer/reparam/__init__.py
--- a/pyro/infer/reparam/__init__.py
+++ b/pyro/infer/reparam/__init__.py
@@ -6,6 +6,7 @@
 from .loc_scale import LocScaleReparam
 from .neutra import NeuTraReparam
 from .stable import LatentStableReparam, StableReparam, SymmetricStableReparam
+from .studentt import StudentTReparam
 from .transform import TransformReparam
 
 __all__ = [
@@ -15,6 +16,7 @@
     "LocScaleReparam",
     "NeuTraReparam",
     "StableReparam",
+    "StudentTReparam",
     "SymmetricStableReparam",
     "TransformReparam",
 ]
diff --git a/pyro/infer/reparam/hmm.py b/pyro/infer/reparam/hmm.py
--- a/pyro/infer/reparam/hmm.py
+++ b/pyro/infer/reparam/hmm.py
@@ -25,6 +25,7 @@ class LinearHMMReparam(Reparam):
     perform inference in the presence of non-Gaussian factors such as
     :meth:`~pyro.distributions.Stable`, :meth:`~pyro.distributions.StudentT` or
     :meth:`~pyro.distributions.LogNormal` , configure with
+    :class:`~pyro.infer.reparam.studentt.StudentTReparam` ,
     :class:`~pyro.infer.reparam.stable.StableReparam` ,
     :class:`~pyro.infer.reparam.stable.SymmetricStableReparam` , etc.  component
     reparameterizers for ``init``, ``trans``, and ``scale``. For example::
diff --git a/pyro/infer/reparam/studentt.py b/pyro/infer/reparam/studentt.py
new file mode 100644
--- /dev/null
+++ b/pyro/infer/reparam/studentt.py
@@ -0,0 +1,37 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+import pyro
+import pyro.distributions as dist
+
+from .reparam import Reparam
+
+
+class StudentTReparam(Reparam):
+    """
+    Auxiliary variable reparameterizer for
+    :class:`~pyro.distributions.StudentT` random variables.
+
+    This is useful in combination with
+    :class:`~pyro.infer.reparam.hmm.LinearHMMReparam` because it allows
+    StudentT processes to be treated as conditionally Gaussian processes,
+    permitting cheap inference via :class:`~pyro.distributions.GaussianHMM` .
+
+    This reparameterizes a :class:`~pyro.distributions.StudentT` by introducing
+    an auxiliary :class:`~pyro.distributions.Gamma` variable conditioned on
+    which the result is :class:`~pyro.distributions.Normal` .
+    """
+    def __call__(self, name, fn, obs):
+        fn, event_dim = self._unwrap(fn)
+        assert isinstance(fn, dist.StudentT)
+
+        # Draw a sample that depends only on df.
+        half_df = fn.df * 0.5
+        gamma = pyro.sample("{}_gamma".format(name),
+                            self._wrap(dist.Gamma(half_df, half_df), event_dim))
+
+        # Construct a scaled Normal.
+        loc = fn.loc
+        scale = fn.scale * gamma.rsqrt()
+        new_fn = self._wrap(dist.Normal(loc, scale), event_dim)
+        return new_fn, obs


Test patch:
diff --git a/tests/infer/reparam/test_hmm.py b/tests/infer/reparam/test_hmm.py
--- a/tests/infer/reparam/test_hmm.py
+++ b/tests/infer/reparam/test_hmm.py
@@ -7,11 +7,18 @@
 import pyro
 import pyro.distributions as dist
 from pyro import poutine
-from pyro.infer.reparam import LinearHMMReparam, StableReparam, SymmetricStableReparam
+from pyro.infer.reparam import LinearHMMReparam, StableReparam, StudentTReparam, SymmetricStableReparam
 from tests.common import assert_close
 from tests.ops.gaussian import random_mvn
 
 
+def random_studentt(shape):
+    df = torch.rand(shape).exp()
+    loc = torch.randn(shape)
+    scale = torch.rand(shape).exp()
+    return dist.StudentT(df, loc, scale)
+
+
 def random_stable(shape, stability, skew=None):
     if skew is None:
         skew = dist.Uniform(-1, 1).sample(shape)
@@ -47,6 +54,31 @@ def model(data=None):
     tr.trace.compute_log_prob()  # smoke test only
 
 
+@pytest.mark.parametrize("duration", [1, 2, 3, 4, 5, 6])
+@pytest.mark.parametrize("obs_dim", [1, 2])
+@pytest.mark.parametrize("hidden_dim", [1, 3])
+@pytest.mark.parametrize("batch_shape", [(), (4,), (2, 3)], ids=str)
+def test_studentt_hmm_shape(batch_shape, duration, hidden_dim, obs_dim):
+    init_dist = random_studentt(batch_shape + (hidden_dim,)).to_event(1)
+    trans_mat = torch.randn(batch_shape + (duration, hidden_dim, hidden_dim))
+    trans_dist = random_studentt(batch_shape + (duration, hidden_dim)).to_event(1)
+    obs_mat = torch.randn(batch_shape + (duration, hidden_dim, obs_dim))
+    obs_dist = random_studentt(batch_shape + (duration, obs_dim)).to_event(1)
+    hmm = dist.LinearHMM(init_dist, trans_mat, trans_dist, obs_mat, obs_dist)
+
+    def model(data=None):
+        with pyro.plate_stack("plates", batch_shape):
+            return pyro.sample("x", hmm, obs=data)
+
+    data = model()
+    rep = StudentTReparam()
+    with poutine.trace() as tr:
+        with poutine.reparam(config={"x": LinearHMMReparam(rep, rep, rep)}):
+            model(data)
+    assert isinstance(tr.trace.nodes["x"]["fn"], dist.GaussianHMM)
+    tr.trace.compute_log_prob()  # smoke test only
+
+
 @pytest.mark.parametrize("duration", [1, 2, 3, 4, 5, 6])
 @pytest.mark.parametrize("obs_dim", [1, 2])
 @pytest.mark.parametrize("hidden_dim", [1, 3])
diff --git a/tests/infer/reparam/test_studentt.py b/tests/infer/reparam/test_studentt.py
new file mode 100644
--- /dev/null
+++ b/tests/infer/reparam/test_studentt.py
@@ -0,0 +1,67 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+from scipy.stats import ks_2samp
+from torch.autograd import grad
+
+import pyro
+import pyro.distributions as dist
+from pyro import poutine
+from pyro.infer.reparam import StudentTReparam
+from tests.common import assert_close
+
+
+# Test helper to extract a few absolute moments from univariate samples.
+# This uses abs moments because StudentT variance may be infinite.
+def get_moments(x):
+    points = torch.tensor([-4., -1., 0., 1., 4.])
+    points = points.reshape((-1,) + (1,) * x.dim())
+    return torch.cat([x.mean(0, keepdim=True), (x - points).abs().mean(1)])
+
+
+@pytest.mark.parametrize("shape", [(), (4,), (2, 3)], ids=str)
+def test_moments(shape):
+    df = torch.empty(shape).uniform_(1.8, 5).requires_grad_()
+    loc = torch.empty(shape).uniform_(-1., 1.).requires_grad_()
+    scale = torch.empty(shape).uniform_(0.5, 1.0).requires_grad_()
+    params = [df, loc, scale]
+
+    def model():
+        with pyro.plate_stack("plates", shape):
+            with pyro.plate("particles", 100000):
+                return pyro.sample("x", dist.StudentT(df, loc, scale))
+
+    value = model()
+    expected_moments = get_moments(value)
+
+    reparam_model = poutine.reparam(model, {"x": StudentTReparam()})
+    trace = poutine.trace(reparam_model).get_trace()
+    assert isinstance(trace.nodes["x"]["fn"], dist.Normal)
+    trace.compute_log_prob()  # smoke test only
+    value = trace.nodes["x"]["value"]
+    actual_moments = get_moments(value)
+    assert_close(actual_moments, expected_moments, atol=0.05)
+
+    for actual_m, expected_m in zip(actual_moments, expected_moments):
+        expected_grads = grad(expected_m.sum(), params, retain_graph=True)
+        actual_grads = grad(actual_m.sum(), params, retain_graph=True)
+        assert_close(actual_grads[0], expected_grads[0], atol=0.2)
+        assert_close(actual_grads[1], expected_grads[1], atol=0.1)
+        assert_close(actual_grads[2], expected_grads[2], atol=0.1)
+
+
+@pytest.mark.parametrize("df", [0.5, 1.0, 1.5, 2.0, 3.0])
+@pytest.mark.parametrize("scale", [0.1, 1.0, 2.0])
+@pytest.mark.parametrize("loc", [0.0, 1.234])
+def test_distribution(df, loc, scale):
+
+    def model():
+        with pyro.plate("particles", 20000):
+            return pyro.sample("x", dist.StudentT(df, loc, scale))
+
+    expected = model()
+    with poutine.reparam(config={"x": StudentTReparam()}):
+        actual = model()
+    assert ks_2samp(expected, actual).pvalue > 0.05