Repository: pyro-ppl/pyro

Issue:
[bug] Continuous autoguides do not take device into account
### Issue Description
Many of the `AutoContinuous` guides do not take into account that the model might be on the GPU instead of the CPU. For example, see the code snippet below, where it fails if I pass `device='cuda'` to the tensors I create.

It would of course be nice if it instead relied on the device given from the model trace.
Do you have any suggestions?

Thanks!

### Environment
 - OS and python version: Windows 10, Python 3.7
 - PyTorch version: 1.1
 - Pyro version: 0.3.3

### Code Snippet
```python
import pyro
import pyro.distributions as dist
import torch
from pyro.contrib.autoguide import AutoDiagonalNormal
from pyro.infer import SVI, Trace_ELBO
from pyro.optim import Adam


def model(xs, prior_loc, prior_scale):
    mu = pyro.sample('mu', dist.Normal(prior_loc, prior_scale))
    return pyro.sample('obs', dist.Normal(mu, prior_scale), obs=xs)


def train():
    xs = torch.tensor([5, 5.5, 5.6, 5.8, 6.3], device='cuda')
    prior_loc = torch.tensor(0.0, device='cuda')
    prior_scale = torch.tensor(1.0, device='cuda')
    guide = AutoDiagonalNormal(model)
    optim = Adam({'lr': 1e-3})
    loss = Trace_ELBO()
    svi = SVI(model, guide, optim, loss)
    svi.step(xs, prior_loc, prior_scale)

if __name__ == "__main__":
  train()
```

Error:
```
Traceback (most recent call last):
  File "C:\tools\miniconda3\envs\scanner-pyro\lib\site-packages\IPython\core\interactiveshell.py", line 3265, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-6-68d49e4377f9>", line 26, in <module>
    train()
  File "<ipython-input-6-68d49e4377f9>", line 23, in train
    svi.step(xs, prior_loc, prior_scale)
  File "C:\tools\miniconda3\envs\scanner-pyro\lib\site-packages\pyro\infer\svi.py", line 99, in step
    loss = self.loss_and_grads(self.model, self.guide, *args, **kwargs)
  File "C:\tools\miniconda3\envs\scanner-pyro\lib\site-packages\pyro\infer\trace_elbo.py", line 136, in loss_and_grads
    surrogate_loss_particle.backward(retain_graph=self.retain_graph)
  File "C:\tools\miniconda3\envs\scanner-pyro\lib\site-packages\torch\tensor.py", line 107, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "C:\tools\miniconda3\envs\scanner-pyro\lib\site-packages\torch\autograd\__init__.py", line 93, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function AddBackward0 returned an invalid gradient at index 1 - expected type torch.cuda.FloatTensor but got torch.FloatTensor
```


Discussion:

Resolving patch:
diff --git a/pyro/contrib/autoguide/__init__.py b/pyro/contrib/autoguide/__init__.py
--- a/pyro/contrib/autoguide/__init__.py
+++ b/pyro/contrib/autoguide/__init__.py
@@ -24,8 +24,10 @@ def model():
 import pyro
 import pyro.distributions as dist
 import pyro.poutine as poutine
+from pyro.contrib.autoguide.initialization import (InitMessenger, init_to_feasible, init_to_mean, init_to_median,
+                                                   init_to_sample)
 from pyro.contrib.util import hessian
-from pyro.distributions.util import broadcast_shape, sum_rightmost
+from pyro.distributions.util import broadcast_shape, eye_like, sum_rightmost
 from pyro.infer.enum import config_enumerate
 from pyro.nn import AutoRegressiveNN
 from pyro.poutine.util import prune_subsample_sites
@@ -47,6 +49,11 @@ def model():
     'AutoLaplaceApproximation',
     'AutoLowRankMultivariateNormal',
     'AutoMultivariateNormal',
+    'init_to_feasible',
+    'init_to_mean',
+    'init_to_median',
+    'init_to_sample',
+    'mean_field_guide_entropy',
 ]
 
 
@@ -267,7 +274,14 @@ class AutoDelta(AutoGuide):
         pyro.param("auto_level", torch.tensor([-1., 0., 1.]))
         pyro.param("auto_concentration", torch.ones(k),
                    constraint=constraints.positive)
+
+    :param callable model: A Pyro model.
+    :param callable init_loc_fn: A per-site initialization function.
+        See :ref:`autoguide-initialization` section for available functions.
     """
+    def __init__(self, model, prefix="auto", init_loc_fn=init_to_median):
+        model = InitMessenger(init_loc_fn)(model)
+        super(AutoDelta, self).__init__(model, prefix=prefix)
 
     def __call__(self, *args, **kwargs):
         """
@@ -287,7 +301,8 @@ def __call__(self, *args, **kwargs):
                 for frame in site["cond_indep_stack"]:
                     if frame.vectorized:
                         stack.enter_context(plates[frame.name])
-                value = pyro.param("{}_{}".format(self.prefix, name), site["value"].detach(),
+                value = pyro.param("{}_{}".format(self.prefix, name),
+                                   site["value"].detach(),
                                    constraint=site["fn"].support)
                 result[name] = pyro.sample(name, dist.Delta(value, event_dim=site["fn"].event_dim))
         return result
@@ -319,7 +334,14 @@ class AutoContinuous(AutoGuide):
     [1] `Automatic Differentiation Variational Inference`,
         Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, David M.
         Blei
+
+    :param callable model: A Pyro model.
+    :param callable init_loc_fn: A per-site initialization function.
+        See :ref:`autoguide-initialization` section for available functions.
     """
+    def __init__(self, model, prefix="auto", init_loc_fn=init_to_median):
+        model = InitMessenger(init_loc_fn)(model)
+        super(AutoContinuous, self).__init__(model, prefix=prefix)
 
     def _setup_prototype(self, *args, **kwargs):
         super(AutoContinuous, self)._setup_prototype(*args, **kwargs)
@@ -337,6 +359,19 @@ def _setup_prototype(self, *args, **kwargs):
         if self.latent_dim == 0:
             raise RuntimeError('{} found no latent variables; Use an empty guide instead'.format(type(self).__name__))
 
+    def _init_loc(self):
+        """
+        Creates an initial latent vector using a per-site init function.
+        """
+        parts = []
+        for name, site in self.prototype_trace.iter_stochastic_nodes():
+            constrained_value = site["value"].detach()
+            unconstrained_value = biject_to(site["fn"].support).inv(constrained_value)
+            parts.append(unconstrained_value.reshape(-1))
+        latent = torch.cat(parts)
+        assert latent.size() == (self.latent_dim,)
+        return latent
+
     def get_posterior(self, *args, **kwargs):
         """
         Returns the posterior distribution.
@@ -466,10 +501,9 @@ def get_posterior(self, *args, **kwargs):
         """
         Returns a MultivariateNormal posterior distribution.
         """
-        loc = pyro.param("{}_loc".format(self.prefix),
-                         lambda: torch.zeros(self.latent_dim))
+        loc = pyro.param("{}_loc".format(self.prefix), self._init_loc)
         scale_tril = pyro.param("{}_scale_tril".format(self.prefix),
-                                lambda: torch.eye(self.latent_dim),
+                                lambda: eye_like(loc, self.latent_dim),
                                 constraint=constraints.lower_cholesky)
         return dist.MultivariateNormal(loc, scale_tril=scale_tril)
 
@@ -504,10 +538,9 @@ def get_posterior(self, *args, **kwargs):
         """
         Returns a diagonal Normal posterior distribution.
         """
-        loc = pyro.param("{}_loc".format(self.prefix),
-                         lambda: torch.zeros(self.latent_dim))
+        loc = pyro.param("{}_loc".format(self.prefix), self._init_loc)
         scale = pyro.param("{}_scale".format(self.prefix),
-                           lambda: torch.ones(self.latent_dim),
+                           lambda: loc.new_ones(self.latent_dim),
                            constraint=constraints.positive)
         return dist.Normal(loc, scale).to_event(1)
 
@@ -542,25 +575,27 @@ class AutoLowRankMultivariateNormal(AutoContinuous):
 
     :param callable model: a generative model
     :param int rank: the rank of the low-rank part of the covariance matrix
+    :param callable init_loc_fn: A per-site initialization function.
+        See :ref:`autoguide-initialization` section for available functions.
     :param str prefix: a prefix that will be prefixed to all param internal sites
     """
 
-    def __init__(self, model, prefix="auto", rank=1):
+    def __init__(self, model, prefix="auto", init_loc_fn=init_to_median, rank=1):
         if not isinstance(rank, numbers.Number) or not rank > 0:
             raise ValueError("Expected rank > 0 but got {}".format(rank))
         self.rank = rank
-        super(AutoLowRankMultivariateNormal, self).__init__(model, prefix)
+        super(AutoLowRankMultivariateNormal, self).__init__(
+            model, prefix=prefix, init_loc_fn=init_loc_fn)
 
     def get_posterior(self, *args, **kwargs):
         """
         Returns a LowRankMultivariateNormal posterior distribution.
         """
-        loc = pyro.param("{}_loc".format(self.prefix),
-                         lambda: torch.zeros(self.latent_dim))
+        loc = pyro.param("{}_loc".format(self.prefix), self._init_loc)
         factor = pyro.param("{}_cov_factor".format(self.prefix),
-                            lambda: torch.randn(self.latent_dim, self.rank) * (0.5 / self.rank) ** 0.5)
+                            lambda: loc.new_empty(self.latent_dim, self.rank).normal_(0, (0.5 / self.rank) ** 0.5))
         diagonal = pyro.param("{}_cov_diag".format(self.prefix),
-                              lambda: torch.ones(self.latent_dim) * 0.5,
+                              lambda: loc.new_full((self.latent_dim,), 0.5),
                               constraint=constraints.positive)
         return dist.LowRankMultivariateNormal(loc, factor, diagonal)
 
@@ -586,13 +621,15 @@ class AutoIAFNormal(AutoContinuous):
 
     :param callable model: a generative model
     :param int hidden_dim: number of hidden dimensions in the IAF
+    :param callable init_loc_fn: A per-site initialization function.
+        See :ref:`autoguide-initialization` section for available functions.
     :param str prefix: a prefix that will be prefixed to all param internal sites
     """
 
-    def __init__(self, model, hidden_dim=None, prefix="auto"):
+    def __init__(self, model, hidden_dim=None, prefix="auto", init_loc_fn=init_to_median):
         self.hidden_dim = hidden_dim
         self.arn = None
-        super(AutoIAFNormal, self).__init__(model, prefix)
+        super(AutoIAFNormal, self).__init__(model, prefix=prefix, init_loc_fn=init_loc_fn)
 
     def get_posterior(self, *args, **kwargs):
         """
@@ -639,8 +676,7 @@ def get_posterior(self, *args, **kwargs):
         """
         Returns a Delta posterior distribution for MAP inference.
         """
-        loc = pyro.param("{}_loc".format(self.prefix),
-                         lambda: torch.zeros(self.latent_dim))
+        loc = pyro.param("{}_loc".format(self.prefix), self._init_loc)
         return dist.Delta(loc).to_event(1)
 
     def laplace_approximation(self, *args, **kwargs):
diff --git a/pyro/contrib/autoguide/initialization.py b/pyro/contrib/autoguide/initialization.py
new file mode 100644
--- /dev/null
+++ b/pyro/contrib/autoguide/initialization.py
@@ -0,0 +1,89 @@
+r"""
+The pyro.contrib.autoguide module contains initialization functions for
+automatic guides.
+
+The standard interface for initialization is a function that inputs a Pyro
+trace ``site`` dict and returns an appropriately sized ``value`` to serve
+as an initial constrained value for a guide estimate.
+"""
+from __future__ import absolute_import, division, print_function
+
+import torch
+from torch.distributions import transform_to
+
+from pyro.poutine.messenger import Messenger
+from pyro.util import torch_isnan
+
+
+def init_to_feasible(site):
+    """
+    Initialize to an arbitrary feasible point, ignoring distribution
+    parameters.
+    """
+    value = site["fn"].sample().detach()
+    t = transform_to(site["fn"].support)
+    return t(torch.zeros_like(t.inv(value)))
+
+
+def init_to_sample(site):
+    """
+    Initialize to a random sample from the prior.
+    """
+    return site["fn"].sample().detach()
+
+
+def init_to_median(site, num_samples=15):
+    """
+    Initialize to the prior median; fallback to a feasible point if median is
+    undefined.
+    """
+    try:
+        # Try to compute empirical median.
+        samples = site["fn"].sample(sample_shape=(num_samples,))
+        value = samples.median(dim=0)[0]
+        if torch_isnan(value):
+            raise ValueError
+        if hasattr(site["fn"], "_validate_sample"):
+            site["fn"]._validate_sample(value)
+        return value
+    except (RuntimeError, ValueError):
+        # Fall back to feasible point.
+        return init_to_feasible(site)
+
+
+def init_to_mean(site):
+    """
+    Initialize to the prior mean; fallback to median if mean is undefined.
+    """
+    try:
+        # Try .mean() method.
+        value = site["fn"].mean.detach()
+        if torch_isnan(value):
+            raise ValueError
+        if hasattr(site["fn"], "_validate_sample"):
+            site["fn"]._validate_sample(value)
+        return value
+    except (NotImplementedError, ValueError):
+        # Fall back to a median.
+        # This is requred for distributions with infinite variance, e.g. Cauchy.
+        return init_to_median(site)
+
+
+class InitMessenger(Messenger):
+    """
+    Initializes a site by replacing ``.sample()`` calls with values
+    drawn from an initialization strategy. This is mainly for internal use by
+    autoguide classes.
+
+    :param callable init_fn: An initialization function.
+    """
+    def __init__(self, init_fn):
+        self.init_fn = init_fn
+        super(InitMessenger, self).__init__()
+
+    def _pyro_sample(self, msg):
+        if msg["done"] or msg["is_observed"] or type(msg["fn"]).__name__ == "_Subsample":
+            return
+        with torch.no_grad():
+            msg["value"] = self.init_fn(msg)
+        msg["done"] = True


Test patch:
diff --git a/tests/contrib/autoguide/test_advi.py b/tests/contrib/autoguide/test_advi.py
--- a/tests/contrib/autoguide/test_advi.py
+++ b/tests/contrib/autoguide/test_advi.py
@@ -1,5 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
+import functools
+
 import numpy as np
 import pytest
 import torch
@@ -10,7 +12,8 @@
 import pyro.poutine as poutine
 from pyro.contrib.autoguide import (AutoCallable, AutoDelta, AutoDiagonalNormal, AutoDiscreteParallel, AutoGuideList,
                                     AutoIAFNormal, AutoLaplaceApproximation, AutoLowRankMultivariateNormal,
-                                    AutoMultivariateNormal)
+                                    AutoMultivariateNormal, init_to_feasible, init_to_mean, init_to_median,
+                                    init_to_sample)
 from pyro.infer import SVI, Trace_ELBO, TraceEnum_ELBO, TraceGraph_ELBO
 from pyro.optim import Adam
 from tests.common import assert_equal
@@ -43,6 +46,12 @@ def model():
 
 
 @pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
+@pytest.mark.parametrize("init_loc_fn", [
+    init_to_feasible,
+    init_to_mean,
+    init_to_median,
+    init_to_sample,
+])
 @pytest.mark.parametrize("auto_class", [
     AutoDelta,
     AutoDiagonalNormal,
@@ -51,7 +60,7 @@ def model():
     AutoIAFNormal,
     AutoLaplaceApproximation,
 ])
-def test_shapes(auto_class, Elbo):
+def test_shapes(auto_class, init_loc_fn, Elbo):
 
     def model():
         pyro.sample("z1", dist.Normal(0.0, 1.0))
@@ -59,7 +68,7 @@ def model():
         with pyro.plate("plate", 3):
             pyro.sample("z3", dist.Normal(torch.zeros(3), torch.ones(3)))
 
-    guide = auto_class(model)
+    guide = auto_class(model, init_loc_fn=init_loc_fn)
     elbo = Elbo(strict_enumeration_warning=False)
     loss = elbo.loss(model, guide)
     assert np.isfinite(loss), loss
@@ -105,7 +114,7 @@ def auto_guide_list_x(model):
 def auto_guide_callable(model):
     def guide_x():
         x_loc = pyro.param("x_loc", torch.tensor(1.))
-        x_scale = pyro.param("x_scale", torch.tensor(2.), constraint=constraints.positive)
+        x_scale = pyro.param("x_scale", torch.tensor(.1), constraint=constraints.positive)
         pyro.sample("x", dist.Normal(x_loc, x_scale))
 
     def median_x():
@@ -125,6 +134,10 @@ def median_x():
     AutoLaplaceApproximation,
     auto_guide_list_x,
     auto_guide_callable,
+    functools.partial(AutoDiagonalNormal, init_loc_fn=init_to_feasible),
+    functools.partial(AutoDiagonalNormal, init_loc_fn=init_to_mean),
+    functools.partial(AutoDiagonalNormal, init_loc_fn=init_to_median),
+    functools.partial(AutoDiagonalNormal, init_loc_fn=init_to_sample),
 ])
 @pytest.mark.parametrize("Elbo", [Trace_ELBO, TraceGraph_ELBO, TraceEnum_ELBO])
 def test_median(auto_class, Elbo):
diff --git a/tests/contrib/autoguide/test_inference.py b/tests/contrib/autoguide/test_inference.py
--- a/tests/contrib/autoguide/test_inference.py
+++ b/tests/contrib/autoguide/test_inference.py
@@ -51,7 +51,7 @@ def do_test_auto(self, N, reparameterized, n_steps):
                      .format(self.target_auto_diag_cov[1:].detach().cpu().numpy()))
 
         # TODO speed up with parallel num_particles > 1
-        adam = optim.Adam({"lr": .0005, "betas": (0.95, 0.999)})
+        adam = optim.Adam({"lr": .001, "betas": (0.95, 0.999)})
         svi = SVI(self.model, self.guide, adam, loss=Trace_ELBO())
 
         for k in range(n_steps):