Repository: pyro-ppl/pyro

Issue:
Create a MockNormal class; drop instance .reparameterized
## Background

Currently `.reparameterized` is a class-level attribute but can be overridden by instance attributes. This override is only used in testing, where we make sure that `TraceGraph_ELBO` performs the same inference on a reparameterized and non-reparameterized `Normal`.

## The problem

The instance attribute `.reparameterized` causes a few problems:
1. It requires extra plumbing for all `Distribution` classes (they must pass `*args,**kwargs` to `super(...).__init__()`)
2. It is not supported by `torch.distributions` to which we are migrating
3. It is not respected by `RandomPrimitive` which simply ignores the instance attribute and returns the class attribute. (This is necessary because `RandomPrimitive` only lazily constructs instances, and there may be no instance to examine).
4. It complicates the torch.distribution wrapper layer. The wrapper dynamically chooses between pyro vs torch implementations (torch does not support fancy things like `batch_size`). Some torch distributions are reparameterized whereas their Pyro equivalents are not (`Gamma`, `Beta`, `Dirichlet`). Therefore `RandomPrimitive` cannot statically determine whether a distribution is reparameterized.

## Proposed solution

We should eliminate the instance-level `.reparameterized` property. For the few cases where we need a non-reparameterized `Normal` in testing, we should simply implement a `NonreparamNormal` class (or even a `Nonreparameterized` class -> class transform).


Discussion:
@martinjankowiak does this sound reasonable? We've discussed this in the past.
yeah sounds good to me. how will TraceGraph_ELBO know whether a given distribution is reparameterized?

also, how does this interact with rsample()?

as long as we have some way to use non-rep normals for testing purposes, i'm pretty happy
> how will TraceGraph_ELBO know whether a given distribution is reparameterized?

There will still be a class property `.reparameterized`. This issues merely suggests disallowing instances to dynamically override that static property.

> how does this interact with `rsample()`?

This is independent of `rsample()`. The `torch.distributions` analog of `.reparameterized` is `.has_rsample`, which is decided per-class rather than per-instance.
ah gotcha

Resolving patch:
diff --git a/pyro/distributions/distribution.py b/pyro/distributions/distribution.py
--- a/pyro/distributions/distribution.py
+++ b/pyro/distributions/distribution.py
@@ -88,17 +88,6 @@ class Distribution(object):
     reparameterized = False
     enumerable = False
 
-    def __init__(self, reparameterized=None):
-        """
-        Constructor for base distribution class.
-
-        :param bool reparameterized: Optional argument to override whether
-            instance should be considered reparameterized (by default, this
-            is decided by the class).
-        """
-        if reparameterized is not None:
-            self.reparameterized = reparameterized
-
     def batch_shape(self, x=None, *args, **kwargs):
         """
         The left-hand tensor shape of samples, used for batching.
diff --git a/pyro/distributions/normal.py b/pyro/distributions/normal.py
--- a/pyro/distributions/normal.py
+++ b/pyro/distributions/normal.py
@@ -153,16 +153,12 @@ def _warn_fallback(message):
 
 @torch_wrapper(Normal)
 def WrapNormal(mu, sigma, batch_size=None, log_pdf_mask=None, *args, **kwargs):
-    reparameterized = kwargs.pop('reparameterized', None)
     if not hasattr(torch, 'distributions'):
         _warn_fallback('Missing module torch.distribution')
     elif not hasattr(torch.distributions, 'Normal'):
         _warn_fallback('Missing class torch.distribution.Normal')
     elif batch_size is not None or log_pdf_mask is not None or args or kwargs:
         _warn_fallback('Unsupported args')
-    elif reparameterized and not TorchNormal.reparameterized:
-        _warn_fallback('Unsupported reparameterized=True')
     else:
-        return TorchNormal(mu, sigma, reparameterized=reparameterized)
-    return Normal(
-        mu, sigma, batch_size=batch_size, log_pdf_mask=log_pdf_mask, reparameterized=reparameterized, *args, **kwargs)
+        return TorchNormal(mu, sigma)
+    return Normal(mu, sigma, batch_size=batch_size, log_pdf_mask=log_pdf_mask, *args, **kwargs)


Test patch:
diff --git a/tests/fakes.py b/tests/fakes.py
new file mode 100644
--- /dev/null
+++ b/tests/fakes.py
@@ -0,0 +1,11 @@
+from __future__ import absolute_import, division, print_function
+
+from pyro.distributions.normal import Normal
+from pyro.distributions.random_primitive import RandomPrimitive
+
+
+class NonreparameterizedNormal(Normal):
+    reparameterized = False
+
+
+nonreparameterized_normal = RandomPrimitive(NonreparameterizedNormal)
diff --git a/tests/infer/test_gradient.py b/tests/infer/test_gradient.py
--- a/tests/infer/test_gradient.py
+++ b/tests/infer/test_gradient.py
@@ -13,6 +13,7 @@
 from pyro.infer import SVI
 from pyro.optim import Adam
 from pyro.util import ng_ones, ng_zeros
+from tests import fakes
 from tests.common import assert_equal
 
 logger = logging.getLogger(__name__)
@@ -27,13 +28,13 @@ def test_subsample_gradient(trace_graph, reparameterized, subsample):
     subsample_size = 1 if subsample else len(data)
     num_particles = 5000
     precision = 0.333
+    normal = dist.normal if reparameterized else fakes.nonreparameterized_normal
 
     def model():
         with pyro.iarange("data", len(data), subsample_size) as ind:
             x = data[ind]
-            z = pyro.sample("z", dist.Normal(ng_zeros(len(x)), ng_ones(len(x)),
-                                             reparameterized=reparameterized))
-            pyro.observe("x", dist.Normal(z, ng_ones(len(x)), reparameterized=reparameterized), x)
+            z = pyro.sample("z", normal, ng_zeros(len(x)), ng_ones(len(x)))
+            pyro.sample("x", normal, z, ng_ones(len(x)), obs=x)
 
     def guide():
         mu = pyro.param("mu", lambda: Variable(torch.zeros(len(data)), requires_grad=True))
@@ -41,7 +42,7 @@ def guide():
         with pyro.iarange("data", len(data), subsample_size) as ind:
             mu = mu[ind]
             sigma = sigma.expand(subsample_size)
-            pyro.sample("z", dist.Normal(mu, sigma, reparameterized=reparameterized))
+            pyro.sample("z", normal, mu, sigma)
 
     optim = Adam({"lr": 0.1})
     inference = SVI(model, guide, optim, loss="ELBO",
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
@@ -14,6 +14,7 @@
 from pyro.distributions.transformed_distribution import TransformedDistribution
 from pyro.infer.svi import SVI
 from pyro.util import ng_ones, ng_zeros
+from tests import fakes
 from tests.common import assert_equal
 from tests.distributions.test_transformed_distribution import AffineExp
 
@@ -75,7 +76,8 @@ def guide():
                                    self.analytic_log_sig_n.data - 0.14 * torch.ones(2),
                                    requires_grad=True))
             sig_q = torch.exp(log_sig_q)
-            pyro.sample("mu_latent", dist.Normal(mu_q, sig_q, reparameterized=reparameterized))
+            normal = dist.normal if reparameterized else fakes.nonreparameterized_normal
+            pyro.sample("mu_latent", normal, mu_q, sig_q)
             pyro.map_data("aaa", self.data, lambda i, x: None,
                           batch_size=self.batch_size)
 
@@ -369,7 +371,8 @@ 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)
-            pyro.sample("mu_latent", dist.Normal(mu_q, sigma, reparameterized=reparameterized))
+            normal = dist.normal if reparameterized else fakes.nonreparameterized_normal
+            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)
diff --git a/tests/integration_tests/test_conjugate_gaussian_models.py b/tests/integration_tests/test_conjugate_gaussian_models.py
--- a/tests/integration_tests/test_conjugate_gaussian_models.py
+++ b/tests/integration_tests/test_conjugate_gaussian_models.py
@@ -15,6 +15,7 @@
 import pyro.distributions as dist
 import pyro.optim as optim
 from pyro.infer import SVI
+from tests import fakes
 from tests.common import assert_equal
 
 logger = logging.getLogger(__name__)
@@ -146,9 +147,8 @@ def guide(*args, **kwargs):
                                              requires_grad=True))
                 mean_function = mu_q if k == self.N else kappa_q * previous_sample + mu_q
                 node_flagged = True if self.which_nodes_reparam[k - 1] == 1.0 else False
-                repa = True if reparameterized else node_flagged
-                latent_dist = dist.Normal(mean_function, sig_q, reparameterized=repa)
-                mu_latent = pyro.sample("mu_latent_%d" % k, latent_dist,
+                normal = dist.normal if reparameterized or node_flagged else fakes.nonreparameterized_normal
+                mu_latent = pyro.sample("mu_latent_%d" % k, normal, mean_function, sig_q,
                                         baseline=dict(use_decaying_avg_baseline=True))
                 previous_sample = mu_latent
             return previous_sample
@@ -436,10 +436,8 @@ def guide(*args, **kwargs):
                                                     requires_grad=True))
                     mean_function_node = mean_function_node + kappa_dep * latents_dict[dep]
                 node_flagged = True if self.which_nodes_reparam[i] == 1.0 else False
-                repa = True if reparameterized else node_flagged
-                latent_dist_node = dist.Normal(mean_function_node, torch.exp(log_sig_node),
-                                               reparameterized=repa)
-                latent_node = pyro.sample(node, latent_dist_node,
+                normal = dist.normal if reparameterized or node_flagged else fakes.nonreparameterized_normal
+                latent_node = pyro.sample(node, normal, mean_function_node, torch.exp(log_sig_node),
                                           baseline=dict(use_decaying_avg_baseline=True,
                                                         baseline_beta=0.96))
                 latents_dict[node] = latent_node
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
@@ -16,6 +16,7 @@
 from pyro.distributions.transformed_distribution import TransformedDistribution
 from pyro.infer import SVI
 from pyro.util import ng_ones, ng_zeros
+from tests import fakes
 from tests.common import assert_equal
 from tests.distributions.test_transformed_distribution import AffineExp
 
@@ -63,11 +64,10 @@ def test_elbo_nonreparameterized(self):
     def do_elbo_test(self, reparameterized, n_steps):
         logger.info(" - - - - - DO NORMALNORMAL ELBO TEST  [reparameterized = %s] - - - - - " % reparameterized)
         pyro.clear_param_store()
+        normal = dist.normal if reparameterized else fakes.nonreparameterized_normal
 
         def model():
-            mu_latent = pyro.sample(
-                    "mu_latent",
-                    dist.Normal(self.mu0, torch.pow(self.lam0, -0.5), reparameterized=reparameterized))
+            mu_latent = pyro.sample("mu_latent", normal, self.mu0, torch.pow(self.lam0, -0.5))
             for i, x in enumerate(self.data):
                 pyro.observe("obs_%d" % i, dist.normal, x, mu_latent,
                              torch.pow(self.lam, -0.5))
@@ -80,8 +80,7 @@ def guide():
                                    self.analytic_log_sig_n.data - 0.29 * torch.ones(2),
                                    requires_grad=True))
             sig_q = torch.exp(log_sig_q)
-            mu_latent = pyro.sample("mu_latent",
-                                    dist.Normal(mu_q, sig_q, reparameterized=reparameterized),
+            mu_latent = pyro.sample("mu_latent", normal, mu_q, sig_q,
                                     baseline=dict(use_decaying_avg_baseline=True))
             return mu_latent
 
@@ -142,6 +141,8 @@ def do_elbo_test(self, repa1, repa2, n_steps, prec, lr, use_nn_baseline, use_dec
         logger.info("[reparameterized = %s, %s; nn_baseline = %s, decaying_baseline = %s]" %
                     (repa1, repa2, use_nn_baseline, use_decaying_avg_baseline))
         pyro.clear_param_store()
+        normal1 = dist.normal if repa1 else fakes.nonreparameterized_normal
+        normal2 = dist.normal if repa2 else fakes.nonreparameterized_normal
 
         if use_nn_baseline:
 
@@ -161,12 +162,8 @@ def forward(self, x):
             mu_prime_baseline = None
 
         def model():
-            mu_latent_prime = pyro.sample(
-                    "mu_latent_prime",
-                    dist.Normal(self.mu0, torch.pow(self.lam0, -0.5), reparameterized=repa1))
-            mu_latent = pyro.sample(
-                    "mu_latent",
-                    dist.Normal(mu_latent_prime, torch.pow(self.lam0, -0.5), reparameterized=repa2))
+            mu_latent_prime = pyro.sample("mu_latent_prime", normal1, self.mu0, torch.pow(self.lam0, -0.5))
+            mu_latent = pyro.sample("mu_latent", normal2, mu_latent_prime, torch.pow(self.lam0, -0.5))
             for i, x in enumerate(self.data):
                 pyro.observe("obs_%d" % i, dist.normal, x, mu_latent,
                              torch.pow(self.lam, -0.5))
@@ -187,14 +184,10 @@ def guide():
                                          Variable(-0.5 * torch.log(1.2 * self.lam0.data),
                                                   requires_grad=True))
             sig_q, sig_q_prime = torch.exp(log_sig_q), torch.exp(log_sig_q_prime)
-            mu_latent_dist = dist.Normal(mu_q, sig_q, reparameterized=repa2)
-            mu_latent = pyro.sample("mu_latent", mu_latent_dist,
+            mu_latent = pyro.sample("mu_latent", normal2, mu_q, sig_q,
                                     baseline=dict(use_decaying_avg_baseline=use_decaying_avg_baseline))
-            mu_latent_prime_dist = dist.Normal(kappa_q.expand_as(mu_latent) * mu_latent + mu_q_prime,
-                                               sig_q_prime,
-                                               reparameterized=repa1)
             pyro.sample("mu_latent_prime",
-                        mu_latent_prime_dist,
+                        normal1, kappa_q.expand_as(mu_latent) * mu_latent + mu_q_prime, sig_q_prime,
                         baseline=dict(nn_baseline=mu_prime_baseline,
                                       nn_baseline_input=mu_latent,
                                       use_decaying_avg_baseline=use_decaying_avg_baseline))
@@ -436,12 +429,12 @@ def test_elbo_nonreparameterized(self):
     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.nonreparameterized_normal
         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))
+            mu_latent = pyro.sample("mu_latent", dist.normal, self.mu0, torch.pow(self.tau0, -0.5))
             sigma = torch.pow(self.tau, -0.5)
             pyro.observe("obs0", dist.lognormal, self.data[0], mu_latent, sigma)
             pyro.observe("obs1", dist.lognormal, self.data[1], mu_latent, sigma)
@@ -451,8 +444,7 @@ 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)
-            pyro.sample("mu_latent",
-                        dist.Normal(mu_q, sigma, reparameterized=reparameterized),
+            pyro.sample("mu_latent", normal, mu_q, sigma,
                         baseline=dict(use_decaying_avg_baseline=True))
 
         adam = optim.Adam({"lr": lr, "betas": (beta1, 0.999)})
@@ -540,9 +532,7 @@ def test_nested_list_map_data_in_elbo(self, n_steps=4000):
         pyro.clear_param_store()
 
         def model():
-            mu_latent = pyro.sample(
-                    "mu_latent",
-                    dist.Normal(self.mu0, torch.pow(self.lam0, -0.5), reparameterized=False))
+            mu_latent = pyro.sample("mu_latent", fakes.nonreparameterized_normal, self.mu0, torch.pow(self.lam0, -0.5))
 
             def obs_outer(i, x):
                 pyro.map_data("map_obs_inner_%d" % i, x, lambda _i, _x:
@@ -564,10 +554,8 @@ def guide():
                                    self.analytic_log_sig_n.data - 0.27 * torch.ones(2),
                                    requires_grad=True))
             sig_q = torch.exp(log_sig_q)
-            mu_latent = pyro.sample(
-                    "mu_latent",
-                    dist.Normal(mu_q, sig_q, reparameterized=False),
-                    baseline=dict(use_decaying_avg_baseline=True))
+            mu_latent = pyro.sample("mu_latent", fakes.nonreparameterized_normal, mu_q, sig_q,
+                                    baseline=dict(use_decaying_avg_baseline=True))
 
             def obs_outer(i, x):
                 pyro.map_data("map_obs_inner_%d" % i, x, lambda _i, _x:
@@ -614,18 +602,16 @@ def _test_vectorized_map_data_in_elbo(self, n_superfluous_top, n_superfluous_bot
                 self.data_tensor[3 * _out + _in, :] = self.data[_out][_in]
 
         def model():
-            mu_latent = pyro.sample(
-                    "mu_latent",
-                    dist.Normal(self.mu0, torch.pow(self.lam0, -0.5), reparameterized=False))
+            mu_latent = pyro.sample("mu_latent", fakes.nonreparameterized_normal, self.mu0, torch.pow(self.lam0, -0.5))
 
             def obs_inner(i, _i, _x):
                 for k in range(n_superfluous_top):
                     pyro.sample("z_%d_%d" % (i, k),
-                                dist.Normal(ng_zeros(4 - i, 1), ng_ones(4 - i, 1), reparameterized=False))
+                                fakes.nonreparameterized_normal, ng_zeros(4 - i, 1), ng_ones(4 - i, 1))
                 pyro.observe("obs_%d" % i, dist.normal, _x, mu_latent, torch.pow(self.lam, -0.5))
                 for k in range(n_superfluous_top, n_superfluous_top + n_superfluous_bottom):
                     pyro.sample("z_%d_%d" % (i, k),
-                                dist.Normal(ng_zeros(4 - i, 1), ng_ones(4 - i, 1), reparameterized=False))
+                                fakes.nonreparameterized_normal, ng_zeros(4 - i, 1), ng_ones(4 - i, 1))
 
             def obs_outer(i, x):
                 pyro.map_data("map_obs_inner_%d" % i, x, lambda _i, _x:
@@ -652,7 +638,7 @@ def guide():
             trivial_baseline = pyro.module("mu_baseline", pt_mu_baseline, tags="baseline")
             baseline_value = trivial_baseline(ng_ones(1))
             mu_latent = pyro.sample("mu_latent",
-                                    dist.Normal(mu_q, sig_q, reparameterized=False),
+                                    fakes.nonreparameterized_normal, mu_q, sig_q,
                                     baseline=dict(baseline_value=baseline_value))
 
             def obs_inner(i, _i, _x):
@@ -663,7 +649,7 @@ def obs_inner(i, _i, _x):
                     mean_i = pyro.param("mean_%d_%d" % (i, k),
                                         Variable(0.5 * torch.ones(4 - i, 1), requires_grad=True))
                     pyro.sample("z_%d_%d" % (i, k),
-                                dist.Normal(mean_i, ng_ones(4 - i, 1), reparameterized=False),
+                                fakes.nonreparameterized_normal, mean_i, ng_ones(4 - i, 1),
                                 baseline=dict(baseline_value=baseline_value))
 
             def obs_outer(i, x):