Repository: pyro-ppl/pyro

Issue:
NotImplementedError in examples/ss_vae_M2.py
Hi,

I had no error before, but after recent update I got an error when run ss_vae_M2.py as:
```
$ python ss_vae_M2.py -enum
Traceback (most recent call last):
  File "ss_vae_M2.py", line 464, in <module>
    run_inference_ss_vae(args)
  File "ss_vae_M2.py", line 371, in run_inference_ss_vae
    run_inference_for_epoch(data_loaders, losses, periodic_interval_batches)
  File "ss_vae_M2.py", line 268, in run_inference_for_epoch
    new_loss = losses[loss_id].step(xs)
  File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/infer/svi.py", line 98, in step
    loss = self.loss_and_grads(self.model, self.guide, *args, **kwargs)
  File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/infer/elbo.py", line 65, in loss_and_grads
    return self.which_elbo.loss_and_grads(model, guide, *args, **kwargs)
  File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/infer/trace_elbo.py", line 146, in loss_and_grads
    for weight, model_trace, guide_trace, log_r in self._get_traces(model, guide, *args, **kwargs):
  File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/infer/trace_elbo.py", line 79, in _get_traces
    check_enum_discrete_can_run(model_trace, guide_trace)
  File "/home/jbaik/.pyenv/versions/3.6.4/lib/python3.6/site-packages/pyro_ppl-0.1.2-py3.6.egg/pyro/infer/trace_elbo.py", line 45, in check_enum_discrete_can_run
    for shape, (source, name) in sorted(shapes.items())])))
NotImplementedError: enum_discrete does not support mixture of batched and un-batched variables. Try rewriting your model to avoid batching or running with enum_discrete=False. Found the following variables of different batch shapes:
guide y: shape = (200,)
guide z: shape = (200, 50)
model x: shape = (200, 784)
```
with `-enum` option. What to do to avoid this?



Discussion:
Hi, what version of PyTorch are you using?  Did you install Pyro from pip or from GitHub?
I'm using Pytorch `0.4.0a0+5b43c22` and Pyro `0.1.2+4d4ffbc`. I've installed them both from each repo.

Pre-release Pyro and PyTorch are both somewhat unstable at the moment due to the ongoing work on `torch.distributions`.  This looks like a bug in `SVI` caused by the new broadcasting semantics in `torch.distributions` (cc @fritzo).  For now, I'd suggest using the `master` branch of Pyro and the current release of PyTorch.

Resolving patch:
diff --git a/examples/ss_vae_M2.py b/examples/ss_vae_M2.py
--- a/examples/ss_vae_M2.py
+++ b/examples/ss_vae_M2.py
@@ -1,27 +1,17 @@
 
 import torch
 import pyro
-import sys
 from torch.autograd import Variable
 import pyro.distributions as dist
 from utils.mnist_cached import MNISTCached, setup_data_loaders
 from pyro.infer import SVI
 from pyro.optim import Adam
 from pyro.nn import ClippedSoftmax, ClippedSigmoid
-from pyro.shim import parse_torch_version
 from utils.custom_mlp import MLP, Exp
 from utils.vae_plots import plot_conditional_samples_ssvae, mnist_test_tsne_ssvae
 from util import set_seed, print_and_log, mkdir_p
 import torch.nn as nn
 
-version_warning = '''
-11/02/2017: This example does not work with PyTorch 0.2, please install PyTorch 0.3.
-'''
-torch_version = parse_torch_version()
-if (torch_version < (0, 2, 1) and not torch_version[-1].startswith("+")):
-    print(version_warning)
-    sys.exit(0)
-
 
 class SSVAE(nn.Module):
     """
@@ -119,7 +109,7 @@ def model(self, xs, ys=None):
             # sample the handwriting style from the constant prior distribution
             prior_mu = Variable(torch.zeros([batch_size, self.z_dim]))
             prior_sigma = Variable(torch.ones([batch_size, self.z_dim]))
-            zs = pyro.sample("z", dist.normal, prior_mu, prior_sigma)
+            zs = pyro.sample("z", dist.normal, prior_mu, prior_sigma, extra_event_dims=1)
 
             # if the label y (which digit to write) is supervised, sample from the
             # constant prior, otherwise, observe the value (i.e. score it against the constant prior)
@@ -134,7 +124,7 @@ def model(self, xs, ys=None):
             # parametrized distribution p(x|y,z) = bernoulli(decoder(y,z))
             # where `decoder` is a neural network
             mu = self.decoder.forward([zs, ys])
-            pyro.sample("x", dist.bernoulli, mu, obs=xs)
+            pyro.sample("x", dist.bernoulli, mu, extra_event_dims=1, obs=xs)
 
     def guide(self, xs, ys=None):
         """
@@ -161,7 +151,7 @@ def guide(self, xs, ys=None):
             # sample (and score) the latent handwriting-style with the variational
             # distribution q(z|x,y) = normal(mu(x,y),sigma(x,y))
             mu, sigma = self.encoder_z.forward([xs, ys])
-            zs = pyro.sample("z", dist.normal, mu, sigma)   # noqa: F841
+            pyro.sample("z", dist.normal, mu, sigma, extra_event_dims=1)
 
     def classifier(self, xs):
         """
@@ -214,11 +204,11 @@ def model_sample(self, ys, batch_size=1):
         # sample the handwriting style from the constant prior distribution
         prior_mu = Variable(torch.zeros([batch_size, self.z_dim]))
         prior_sigma = Variable(torch.ones([batch_size, self.z_dim]))
-        zs = pyro.sample("z", dist.normal, prior_mu, prior_sigma)
+        zs = pyro.sample("z", dist.normal, prior_mu, prior_sigma, extra_event_dims=1)
 
         # sample an image using the decoder
         mu = self.decoder.forward([zs, ys])
-        xs = pyro.sample("sample", dist.bernoulli, mu)
+        xs = pyro.sample("sample", dist.bernoulli, mu, extra_event_dims=1)
         return xs, mu
 
 


Test patch:
diff --git a/tests/test_examples.py b/tests/test_examples.py
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -2,15 +2,14 @@
 
 import os
 import sys
-from collections import OrderedDict
 from subprocess import check_call
 
 import pytest
 
 from tests.common import EXAMPLES_DIR, requires_cuda
 
-CPU_EXAMPLES = OrderedDict()
-CUDA_EXAMPLES = OrderedDict()
+CPU_EXAMPLES = []
+CUDA_EXAMPLES = []
 
 
 def discover_examples():
@@ -31,16 +30,24 @@ def discover_examples():
                 # Either this is not a main file, or we don't know how to run it cheaply.
                 continue
             example = os.path.relpath(path, EXAMPLES_DIR)
-            CPU_EXAMPLES[example] = args
+            CPU_EXAMPLES.append((example, args))
+            if '--enum-discrete' in text:
+                CPU_EXAMPLES.append((example, args + ['--enum-discrete']))
             if '--cuda' in text:
-                CUDA_EXAMPLES[example] = args + ['--cuda']
+                CUDA_EXAMPLES.append((example, args + ['--cuda']))
+    CPU_EXAMPLES.sort()
+    CUDA_EXAMPLES.sort()
 
 
 discover_examples()
 
 
+def make_ids(examples):
+    return ['{} {}'.format(example, ' '.join(args)) for example, args in examples]
+
+
 @pytest.mark.stage("test_examples")
-@pytest.mark.parametrize('example,args', CPU_EXAMPLES.items(), ids=list(CPU_EXAMPLES))
+@pytest.mark.parametrize('example,args', CPU_EXAMPLES, ids=make_ids(CPU_EXAMPLES))
 def test_cpu(example, args):
     example = os.path.join(EXAMPLES_DIR, example)
     check_call([sys.executable, example] + args)
@@ -48,7 +55,7 @@ def test_cpu(example, args):
 
 @requires_cuda
 @pytest.mark.stage("test_examples")
-@pytest.mark.parametrize('example,args', CUDA_EXAMPLES.items(), ids=list(CUDA_EXAMPLES))
+@pytest.mark.parametrize('example,args', CUDA_EXAMPLES, ids=make_ids(CUDA_EXAMPLES))
 def test_cuda(example, args):
     example = os.path.join(EXAMPLES_DIR, example)
     check_call([sys.executable, example] + args)