# Buggy code
def guide(self, xs, ys=None):
    """
    The guide corresponds to the following:
    q(y|x) = categorical(alpha(x))              # infer digit from an image
    q(z|x,y) = normal(mu(x,y),sigma(x,y))       # infer handwriting style from an image and the digit
    mu, sigma are given by a neural network `encoder_z`
    alpha is given by a neural network `encoder_y`
    :param xs: a batch of scaled vectors of pixels from an image
    :param ys: (optional) a batch of the class labels i.e.
               the digit corresponding to the image(s)
    :return: None
    """
    # inform Pyro that the variables in the batch of xs, ys are conditionally independent
    with pyro.iarange("independent"):
        # if the class label (the digit) is not supervised, sample
        # (and score) the digit with the variational distribution
        # q(y|x) = categorical(alpha(x))
        if ys is None:
            alpha = self.encoder_y.forward(xs)
            ys = pyro.sample("y", dist.one_hot_categorical, alpha)

        # 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


# Traceback (NotImplementedError)
$ 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)

# Fixed code
def guide(self, xs, ys=None):
    """
    The guide corresponds to the following:
    q(y|x) = categorical(alpha(x))              # infer digit from an image
    q(z|x,y) = normal(mu(x,y),sigma(x,y))       # infer handwriting style from an image and the digit
    mu, sigma are given by a neural network `encoder_z`
    alpha is given by a neural network `encoder_y`
    :param xs: a batch of scaled vectors of pixels from an image
    :param ys: (optional) a batch of the class labels i.e.
               the digit corresponding to the image(s)
    :return: None
    """
    # inform Pyro that the variables in the batch of xs, ys are conditionally independent
    with pyro.iarange("independent"):
        # if the class label (the digit) is not supervised, sample
        # (and score) the digit with the variational distribution
        # q(y|x) = categorical(alpha(x))
        if ys is None:
            alpha = self.encoder_y.forward(xs)
            ys = pyro.sample("y", dist.one_hot_categorical, alpha)

        # 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])
        pyro.sample("z", dist.normal, mu, sigma, extra_event_dims=1)