Trying to impute missing categorical data

This work around doesn’t seem to work for me.

When I cast to int64, the sampled values aren’t in the range of my categorical variable and hence I get index errors:

data = np.ma.masked_equal([1, 1, 0, 0, 2, -1, -1], -1)
with pm.Model():
    idx = pm.Categorical("idx", p=[0.1, 0.2, 0.7], observed=data)
    idx = at.cast(idx, 'int64')
    num = pm.Normal("num", 0, 1, size=3)[idx]
    pm.sample()

Output:

    rval = inputs[0].__getitem__(tuple(inputs[1:]))
IndexError: index 5 is out of bounds for axis 0 with size 3