VI for mixture models fails with `IndexError`

[ColCarroll]

This looks like a bug when drawing from certain distributions using VI. In particular, the model does fine with MCMC sampling (the sampling does fine anyways, this is a demonstration of label switching). This is installed from git on master, with Python 3.6. Maybe @ferrine has input?

import numpy as np
import pymc3 as pm

# Make reproducible
np.random.seed(1)

# Generate data
μ_true = np.array([-2, 0, 2])
z_true = np.random.randint(len(μ_true), size=100)
y = np.random.normal(μ_true[z_true], np.ones_like(z_true))

with pm.Model():
    μ = pm.Normal('μ', mu=0, sd=10, shape=3)
    z = pm.Categorical('z', p=tt.ones(3) / 3, shape=len(y))
    y_obs = pm.Normal('y_obs', mu=μ[z], sd=1., observed=y)

    mcmc_trace = pm.sample(1000)  # runs fine

    inference = pm.fit(n=30000)  # fails
    vi_trace = inference.sample(5000)

The call to pm.fit fails with

IndexError: index 3 is out of bounds for size 3
Apply node that caused the error: AdvancedSubtensor1(μ_vi_replacement, z_vi_replacement)
Toposort index: 29
Inputs types: [TensorType(float64, vector), TensorType(int64, vector)]
Inputs shapes: [(3,), (100,)]
Inputs strides: [(8,), (8,)]
Inputs values: [array([ 0.49079266, -0.54784519, -0.40646094]), 'not shown']
Outputs clients: [[Elemwise{Sub}[(0, 1)](TensorConstant{[ 0.120158...61838026]}, AdvancedSubtensor1.0)]]

Backtrace when the node is created(use Theano flag traceback.limit=N to make it longer):
  File "/home/colin/miniconda3/envs/scratch3.6/lib/python3.6/site-packages/ipykernel/kernelbase.py", line 233, in dispatch_shell
    handler(stream, idents, msg)
  File "/home/colin/miniconda3/envs/scratch3.6/lib/python3.6/site-packages/ipykernel/kernelbase.py", line 399, in execute_request
    user_expressions, allow_stdin)
  File "/home/colin/miniconda3/envs/scratch3.6/lib/python3.6/site-packages/ipykernel/ipkernel.py", line 208, in do_execute
    res = shell.run_cell(code, store_history=store_history, silent=silent)
  File "/home/colin/miniconda3/envs/scratch3.6/lib/python3.6/site-packages/ipykernel/zmqshell.py", line 537, in run_cell
    return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
  File "/home/colin/miniconda3/envs/scratch3.6/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 2728, in run_cell
    interactivity=interactivity, compiler=compiler, result=result)
  File "/home/colin/miniconda3/envs/scratch3.6/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 2850, in run_ast_nodes
    if self.run_code(code, result):
  File "/home/colin/miniconda3/envs/scratch3.6/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 2910, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-29-3fb2dfb06212>", line 4, in <module>
    y_obs = pm.Normal('y_obs', mu=μ[z], sd=1., observed=y)

HINT: Use the Theano flag 'exception_verbosity=high' for a debugprint and storage map footprint of this apply node.

[junpenglao]

You cant do VI in pymc3 on model with discrete variables.

[ColCarroll]

Thanks! Makes sense.