Repository: pyro-ppl/pyro

Issue:
tracegraph_elbo related cleanup
- [x] refactor vectorized_map_data_info and related data structures
- [x] refactor way prune_subsample_sites etc. is used to setup a trace for computation
- [x] rethink any needed asserts/warnings in tracegraph_elbo or trace_poutine
- [x] go through test_valid_models and make sure we have good coverage
- [x] rethink integration/gradient test coverage (where possibly prefer cheaper tests)



Discussion:

Resolving patch:
diff --git a/pyro/infer/trace_elbo.py b/pyro/infer/trace_elbo.py
--- a/pyro/infer/trace_elbo.py
+++ b/pyro/infer/trace_elbo.py
@@ -26,7 +26,22 @@ def _compute_log_r(model_trace, guide_trace):
 
 class Trace_ELBO(ELBO):
     """
-    A trace implementation of ELBO-based SVI
+    A trace implementation of ELBO-based SVI. The estimator is constructed
+    along the lines of references [1] and [2]. There are no restrictions on the
+    dependency structure of the model or the guide. The gradient estimator includes
+    partial Rao-Blackwellization for reducing the variance of the estimator when
+    non-reparameterizable random variables are present. The Rao-Blackwellization is
+    partial in that it only uses conditional independence information that is marked
+    by :class:`~pyro.iarange` contexts. For more fine-grained Rao-Blackwellization,
+    see :class:`~pyro.infer.tracegraph_elbo.TraceGraph_ELBO`.
+
+    References
+
+    [1] Automated Variational Inference in Probabilistic Programming,
+        David Wingate, Theo Weber
+
+    [2] Black Box Variational Inference,
+        Rajesh Ranganath, Sean Gerrish, David M. Blei
     """
 
     def _get_traces(self, model, guide, *args, **kwargs):
diff --git a/pyro/infer/tracegraph_elbo.py b/pyro/infer/tracegraph_elbo.py
--- a/pyro/infer/tracegraph_elbo.py
+++ b/pyro/infer/tracegraph_elbo.py
@@ -170,9 +170,14 @@ class TraceGraph_ELBO(ELBO):
     A TraceGraph implementation of ELBO-based SVI. The gradient estimator
     is constructed along the lines of reference [1] specialized to the case
     of the ELBO. It supports arbitrary dependency structure for the model
-    and guide as well as baselines for non-reparameteriable random variables.
-    Where possible, dependency information as recorded in the TraceGraph is
-    used to reduce the variance of the gradient estimator.
+    and guide as well as baselines for non-reparameterizable random variables.
+    Where possible, conditional dependency information as recorded in the
+    :class:`~pyro.poutine.trace.Trace` is used to reduce the variance of the gradient estimator.
+    In particular three kinds of conditional dependency information are
+    used to reduce variance:
+    - the sequential order of samples (z is sampled after y => y does not depend on z)
+    - :class:`~pyro.iarange` generators
+    - :class:`~pyro.irange` generators
 
     References
 


Test patch:
diff --git a/tests/distributions/test_transformed_distribution.py b/tests/distributions/test_transformed_distribution.py
deleted file mode 100644
--- a/tests/distributions/test_transformed_distribution.py
+++ /dev/null
@@ -1,119 +0,0 @@
-from __future__ import absolute_import, division, print_function
-
-import numpy as np
-import pytest
-import scipy.stats as sp
-import torch
-from torch.distributions import AffineTransform, ExpTransform, ComposeTransform
-
-import pyro.distributions as dist
-from pyro.distributions import LogNormal
-from pyro.distributions import TransformedDistribution
-from tests.common import assert_equal
-from tests.distributions.dist_fixture import Fixture
-
-
-def make_lognormal(kwargs):
-    kwargs['examples'] = [kwargs.pop('example')]
-    return Fixture(pyro_dist=LogNormal,
-                   scipy_dist=sp.lognorm,
-                   scipy_arg_fn=lambda mu, sigma: ((np.array(sigma),),
-                                                   {"scale": np.exp(np.array(mu))}),
-                   **kwargs)
-
-
-EXAMPLES = list(map(make_lognormal, [
-    {
-        'example': {'mu': [1.4], 'sigma': [0.4], 'test_data': [5.5]},
-    },
-    {
-        'example': {'mu': [1.4], 'sigma': [0.4], 'test_data': [[5.5]]},
-    },
-    {
-        'example': {
-            'mu': [1.4, 0.4, 0.4],
-            'sigma': [1.2, 0.5, 0.5],
-            'test_data': [[5.5, 6.4, 6.4], [0.5, 0.4, 0.4]],
-        },
-        'min_samples': 500000,
-    },
-    {
-        'example': {'mu': [1.4], 'sigma': [1.2], 'test_data': [[5.5], [6.4]]},
-        'min_samples': 1000000,
-    },
-]))
-
-
-def AffineExp(affine_b, affine_a):
-    affine_transform = AffineTransform(loc=affine_a, scale=affine_b)
-    exp_transform = ExpTransform()
-    return ComposeTransform([affine_transform, exp_transform])
-
-
-def get_transformed_dist(distribution, affine_a, affine_b):
-    return TransformedDistribution(distribution, [AffineExp(affine_b, affine_a)])
-
-
-@pytest.mark.parametrize('lognormal', EXAMPLES)
-def test_mean_and_var(lognormal):
-    dist_params = lognormal.get_dist_params(0)
-    mu_lognorm = dist_params['mu']
-    sigma_lognorm = dist_params['sigma']
-    mu_z = torch.zeros_like(mu_lognorm)
-    sigma_z = torch.ones_like(sigma_lognorm)
-    normal_dist = dist.Normal(mu_z, sigma_z)
-    trans_dist = get_transformed_dist(normal_dist, mu_lognorm, sigma_lognorm)
-    torch_samples = trans_dist.sample(sample_shape=torch.Size((lognormal.get_num_samples(0),)))
-    torch_mean = torch.mean(torch_samples, 0)
-    torch_std = torch.std(torch_samples, 0)
-    analytic_mean = lognormal.pyro_dist(**dist_params).mean
-    analytic_std = lognormal.pyro_dist(**dist_params).variance ** 0.5
-    precision = analytic_mean.max().item() * 0.05
-    assert_equal(torch_mean, analytic_mean, prec=precision)
-    assert_equal(torch_std, analytic_std, prec=precision)
-
-
-@pytest.mark.parametrize('lognormal', EXAMPLES)
-def test_log_pdf(lognormal):
-    dist_params = lognormal.get_dist_params(0)
-    mu_lognorm = dist_params['mu']
-    sigma_lognorm = dist_params['sigma']
-    mu_z = torch.zeros_like(mu_lognorm)
-    sigma_z = torch.ones_like(sigma_lognorm)
-    normal_dist = dist.Normal(mu_z, sigma_z)
-    trans_dist = get_transformed_dist(normal_dist, mu_lognorm, sigma_lognorm)
-    test_data = lognormal.get_test_data(0)
-    log_px_torch = trans_dist.log_prob(test_data).sum().item()
-    log_px_np = sp.lognorm.logpdf(
-        test_data.detach().cpu().numpy(),
-        sigma_lognorm.detach().cpu().numpy(),
-        scale=np.exp(mu_lognorm)).sum()
-    assert_equal(log_px_torch, log_px_np, prec=1e-4)
-
-
-@pytest.mark.parametrize('lognormal', EXAMPLES)
-def test_log_prob(lognormal):
-    dist_params = lognormal.get_dist_params(0)
-    mu_lognorm = dist_params['mu']
-    sigma_lognorm = dist_params['sigma']
-    mu_z = torch.zeros_like(mu_lognorm)
-    sigma_z = torch.ones_like(sigma_lognorm)
-    normal_dist = dist.Normal(mu_z, sigma_z)
-    trans_dist = get_transformed_dist(normal_dist, mu_lognorm, sigma_lognorm)
-    test_data = lognormal.get_test_data(0)
-    log_px_torch = trans_dist.log_prob(test_data).detach().cpu().numpy()
-    log_px_np = sp.lognorm.logpdf(
-        test_data.detach().cpu().numpy(),
-        sigma_lognorm.detach().cpu().numpy(),
-        scale=np.exp(mu_lognorm.detach().cpu().numpy()))
-    assert_equal(log_px_torch, log_px_np, prec=1e-4)
-
-
-@pytest.mark.parametrize('lognormal', EXAMPLES)
-def test_shape(lognormal):
-    dist_params = lognormal.get_dist_params(0)
-    mu_lognorm = dist_params['mu']
-    sigma_lognorm = dist_params['sigma']
-    normal_dist = dist.Normal(**dist_params)
-    trans_dist = get_transformed_dist(normal_dist, mu_lognorm, sigma_lognorm)
-    assert_equal(trans_dist.sample().size(), trans_dist.shape())
diff --git a/tests/infer/test_elbo_mapdata.py b/tests/infer/test_elbo_mapdata.py
--- a/tests/infer/test_elbo_mapdata.py
+++ b/tests/infer/test_elbo_mapdata.py
@@ -16,7 +16,7 @@
 
 @pytest.mark.stage("integration", "integration_batch_1")
 @pytest.mark.init(rng_seed=161)
-@pytest.mark.parametrize("batch_size", [3, 5, 7, 8, None])
+@pytest.mark.parametrize("batch_size", [3, 8, None])
 @pytest.mark.parametrize("map_type", ["iarange", "irange", "range"])
 def test_elbo_mapdata(batch_size, map_type):
     # normal-normal: known covariance
diff --git a/tests/infer/test_inference.py b/tests/infer/test_inference.py
--- a/tests/infer/test_inference.py
+++ b/tests/infer/test_inference.py
@@ -5,19 +5,15 @@
 
 import pytest
 import torch
-from torch import nn as nn
 from torch.distributions import constraints
-from torch.nn import Parameter
 
 import pyro
 import pyro.distributions as dist
 import pyro.optim as optim
-from pyro.distributions import TransformedDistribution
 from pyro.distributions.testing import fakes
 from pyro.distributions.testing.rejection_gamma import ShapeAugmentedGamma
 from pyro.infer.svi import SVI
 from tests.common import assert_equal
-from tests.distributions.test_transformed_distribution import AffineExp
 
 
 def param_mse(name, target):
@@ -306,108 +302,6 @@ def guide():
         assert_equal(0.0, beta_error, prec=0.08)
 
 
-class LogNormalNormalGuide(nn.Module):
-    def __init__(self, mu_q_log_init, tau_q_log_init):
-        super(LogNormalNormalGuide, self).__init__()
-        self.mu_q_log = Parameter(mu_q_log_init)
-        self.tau_q_log = Parameter(tau_q_log_init)
-
-
-@pytest.mark.stage("integration", "integration_batch_2")
-class LogNormalNormalTests(TestCase):
-    def setUp(self):
-        # lognormal-normal model
-        # putting some of the parameters inside of a torch module to
-        # make sure that that functionality is ok (XXX: do this somewhere else in the future)
-        self.mu0 = torch.tensor(1.0)  # normal prior hyperparameter
-        # normal prior hyperparameter
-        self.tau0 = torch.tensor(1.0)
-        # known precision for observation likelihood
-        self.tau = torch.tensor(2.5)
-        self.n_data = 2
-        self.data = torch.tensor([1.5, 2.2])  # two observations
-        self.tau_n = self.tau0 + torch.tensor(self.n_data) * self.tau  # posterior tau
-        mu_numerator = self.mu0 * self.tau0 + \
-            self.tau * torch.sum(torch.log(self.data))
-        self.mu_n = mu_numerator / self.tau_n  # posterior mu
-        self.log_mu_n = torch.log(self.mu_n)
-        self.log_tau_n = torch.log(self.tau_n)
-
-    def test_elbo_reparameterized(self):
-        self.do_elbo_test(True, 12000)
-
-    def test_elbo_nonreparameterized(self):
-        self.do_elbo_test(False, 15000)
-
-    def do_elbo_test(self, reparameterized, n_steps):
-        pyro.clear_param_store()
-        pt_guide = LogNormalNormalGuide(self.log_mu_n.data + 0.17,
-                                        self.log_tau_n.data - 0.143)
-
-        def model():
-            mu_latent = pyro.sample("mu_latent",
-                                    dist.Normal(self.mu0, torch.pow(self.tau0, -0.5)))
-            sigma = torch.pow(self.tau, -0.5)
-            with pyro.iarange("iarange", self.n_data):
-                pyro.observe("obs", dist.LogNormal(mu_latent, sigma), obs=self.data)
-            return mu_latent
-
-        def guide():
-            pyro.module("mymodule", pt_guide)
-            mu_q, tau_q = torch.exp(pt_guide.mu_q_log), torch.exp(pt_guide.tau_q_log)
-            sigma = torch.pow(tau_q, -0.5)
-            Normal = dist.Normal if reparameterized else fakes.NonreparameterizedNormal
-            pyro.sample("mu_latent", Normal(mu_q, sigma))
-
-        adam = optim.Adam({"lr": .0005, "betas": (0.96, 0.999)})
-        svi = SVI(model, guide, adam, loss="ELBO", trace_graph=False)
-
-        for k in range(n_steps):
-            svi.step()
-
-        mu_error = param_abs_error("mymodule$$$mu_q_log", self.log_mu_n)
-        tau_error = param_abs_error("mymodule$$$tau_q_log", self.log_tau_n)
-        assert_equal(0.0, mu_error, prec=0.07)
-        assert_equal(0.0, tau_error, prec=0.07)
-
-    def test_elbo_with_transformed_distribution(self):
-        pyro.clear_param_store()
-
-        def model():
-            zero = torch.zeros(1)
-            one = torch.ones(1)
-            mu_latent = pyro.sample("mu_latent",
-                                    dist.Normal(self.mu0, torch.pow(self.tau0, -0.5)))
-            bijector = AffineExp(torch.pow(self.tau, -0.5), mu_latent)
-            x_dist = TransformedDistribution(dist.Normal(zero, one), bijector)
-            with pyro.iarange("data", self.n_data):
-                pyro.observe("obs", x_dist, self.data)
-            return mu_latent
-
-        def guide():
-            mu_q_log = pyro.param(
-                "mu_q_log",
-                torch.tensor(
-                    self.log_mu_n.data +
-                    0.17,
-                    requires_grad=True))
-            tau_q_log = pyro.param("tau_q_log", torch.tensor(self.log_tau_n.data - 0.143,
-                                                             requires_grad=True))
-            mu_q, tau_q = torch.exp(mu_q_log), torch.exp(tau_q_log)
-            pyro.sample("mu_latent", dist.Normal(mu_q, torch.pow(tau_q, -0.5)))
-
-        adam = optim.Adam({"lr": .0005, "betas": (0.96, 0.999)})
-        svi = SVI(model, guide, adam, loss="ELBO", trace_graph=False)
-
-        for k in range(12001):
-            svi.step()
-
-        mu_error = param_abs_error("mu_q_log", self.log_mu_n)
-        tau_error = param_abs_error("tau_q_log", self.log_tau_n)
-        assert_equal(0.0, mu_error, prec=0.05)
-        assert_equal(0.0, tau_error, prec=0.05)
-
-
 class SafetyTests(TestCase):
 
     def setUp(self):
diff --git a/tests/integration_tests/test_tracegraph_elbo.py b/tests/integration_tests/test_tracegraph_elbo.py
--- a/tests/integration_tests/test_tracegraph_elbo.py
+++ b/tests/integration_tests/test_tracegraph_elbo.py
@@ -7,16 +7,13 @@
 import pytest
 import torch
 from torch import nn as nn
-from torch.nn import Parameter
 
 import pyro
 import pyro.distributions as dist
 import pyro.optim as optim
 from pyro.distributions.testing import fakes
-from pyro.distributions import TransformedDistribution
 from pyro.infer import SVI
 from tests.common import assert_equal
-from tests.distributions.test_transformed_distribution import AffineExp
 
 pytestmark = pytest.mark.stage("integration", "integration_batch_2")
 logger = logging.getLogger(__name__)
@@ -329,107 +326,6 @@ def guide():
         assert_equal(0.0, beta_error, prec=0.04)
 
 
-class LogNormalNormalGuide(nn.Module):
-    def __init__(self, mu_q_log_init, tau_q_log_init):
-        super(LogNormalNormalGuide, self).__init__()
-        self.mu_q_log = Parameter(mu_q_log_init)
-        self.tau_q_log = Parameter(tau_q_log_init)
-
-
-class LogNormalNormalTests(TestCase):
-    def setUp(self):
-        # lognormal-normal model
-        # putting some of the parameters inside of a torch module to
-        # make sure that that functionality is ok (XXX: do this somewhere else in the future)
-        self.mu0 = torch.tensor(1.0)  # normal prior hyperparameter
-        # normal prior hyperparameter
-        self.tau0 = torch.tensor(1.0)
-        # known precision for observation likelihood
-        self.tau = torch.tensor(2.5)
-        self.data = torch.tensor([1.5, 2.2])  # two observations
-        self.tau_n = self.tau0 + len(self.data) * self.tau  # posterior tau
-        mu_numerator = self.mu0 * self.tau0 + self.tau * torch.log(self.data).sum(0)
-        self.mu_n = mu_numerator / self.tau_n  # posterior mu
-        self.log_mu_n = torch.log(self.mu_n)
-        self.log_tau_n = torch.log(self.tau_n)
-
-    def test_elbo_reparameterized(self):
-        self.do_elbo_test(True, 7000, 0.95, 0.001)
-
-    def test_elbo_nonreparameterized(self):
-        self.do_elbo_test(False, 7000, 0.95, 0.001)
-
-    def do_elbo_test(self, reparameterized, n_steps, beta1, lr):
-        logger.info(" - - - - - DO LOGNORMAL-NORMAL ELBO TEST [repa = %s] - - - - - " % reparameterized)
-        pyro.clear_param_store()
-        Normal = dist.Normal if reparameterized else fakes.NonreparameterizedNormal
-        pt_guide = LogNormalNormalGuide(self.log_mu_n.data + 0.17,
-                                        self.log_tau_n.data - 0.143)
-
-        def model():
-            mu_latent = pyro.sample("mu_latent", dist.Normal(self.mu0, torch.pow(self.tau0, -0.5)))
-            sigma = torch.pow(self.tau, -0.5)
-            with pyro.iarange("data", len(self.data)):
-                pyro.sample("obs", dist.LogNormal(mu_latent, sigma), obs=self.data)
-            return mu_latent
-
-        def guide():
-            pyro.module("mymodule", pt_guide)
-            mu_q = torch.exp(pt_guide.mu_q_log).squeeze()
-            tau_q = torch.exp(pt_guide.tau_q_log).squeeze()
-            sigma = torch.pow(tau_q, -0.5)
-            pyro.sample("mu_latent", Normal(mu_q, sigma),
-                        infer=dict(baseline=dict(use_decaying_avg_baseline=True)))
-            with pyro.iarange("data", len(self.data)):
-                pass
-
-        adam = optim.Adam({"lr": lr, "betas": (beta1, 0.999)})
-        svi = SVI(model, guide, adam, loss="ELBO", trace_graph=True)
-
-        for k in range(n_steps):
-            svi.step()
-            mu_error = param_abs_error("mymodule$$$mu_q_log", self.log_mu_n)
-            tau_error = param_abs_error("mymodule$$$tau_q_log", self.log_tau_n)
-            if k % 500 == 0:
-                logger.debug("mu_error, tau_error = %.4f, %.4f" % (mu_error, tau_error))
-
-        assert_equal(0.0, mu_error, prec=0.05)
-        assert_equal(0.0, tau_error, prec=0.05)
-
-    def test_elbo_with_transformed_distribution(self):
-        logger.info(" - - - - - DO LOGNORMAL-NORMAL ELBO TEST [uses TransformedDistribution] - - - - - ")
-        pyro.clear_param_store()
-
-        def model():
-            mu_latent = pyro.sample("mu_latent", dist.Normal(self.mu0, torch.pow(self.tau0, -0.5)))
-            bijector = AffineExp(torch.pow(self.tau, -0.5), mu_latent)
-            x_dist = TransformedDistribution(dist.Normal(0, 1), bijector)
-            pyro.sample("obs0", x_dist, obs=self.data[0])
-            pyro.sample("obs1", x_dist, obs=self.data[1])
-            return mu_latent
-
-        def guide():
-            mu_q_log = pyro.param("mu_q_log",
-                                  torch.tensor(self.log_mu_n + 0.17, requires_grad=True))
-            tau_q_log = pyro.param("tau_q_log",
-                                   torch.tensor(self.log_tau_n - 0.143, requires_grad=True))
-            mu_q, tau_q = torch.exp(mu_q_log), torch.exp(tau_q_log)
-            pyro.sample("mu_latent", dist.Normal(mu_q, torch.pow(tau_q, -0.5)))
-
-        adam = optim.Adam({"lr": 0.001, "betas": (0.95, 0.999)})
-        svi = SVI(model, guide, adam, loss="ELBO", trace_graph=True)
-
-        for k in range(7000):
-            svi.step()
-            mu_error = param_abs_error("mu_q_log", self.log_mu_n)
-            tau_error = param_abs_error("tau_q_log", self.log_tau_n)
-            if k % 500 == 0:
-                logger.debug("mu_error, tau_error = %.4f, %.4f" % (mu_error, tau_error))
-
-        assert_equal(0.0, mu_error, prec=0.05)
-        assert_equal(0.0, tau_error, prec=0.05)
-
-
 @pytest.mark.init(rng_seed=0)
 @pytest.mark.stage("integration", "integration_batch_1")
 class RaoBlackwellizationTests(TestCase):