Categorical variable with p equal to deterministic that depends on observed categorical

import pymc3 as pm3
import numpy as np
import theano.tensor as tt
from theano.compile.ops import as_op
arr_y = np.ones(shape=(3, 3, 3))*1/3

# define model
basic_model = pm3.Model()
with basic_model:
    a = pm3.Categorical('a', p=np.array([0.25, .30, .45]), observed = np.array([0, 0]))
    b = pm3.Categorical('b', p=np.array([0.6, .3, .1]), observed = None)

    @as_op(itypes=[tt.lscalar]*2, otypes=[tt.dvector])
    def p_y(a1, b1):
        return arr_y[a1, b1, :]
    y = pm3.Categorical('y', p=p_y(a, b))

This doesn’t work. What am i doing wrong? Personally, I think it’s a bug. I get

TypeError: We expected inputs of types ‘[TensorType(int64, scalar), TensorType(int64, scalar)]’ but got types ‘[TensorType(int64, vector), TensorType(int64, scalar)]’

Transfering arr_y into a theano tensor and index it directly is probably the easiest.

import theano
arr_ys = theano.shared(arr_y)
with pm3.Model():
    a = pm3.Categorical('a', p=np.array([0.25, .30, .45]), observed = np.array([0, 0]))
    b = pm3.Categorical('b', p=np.array([0.6, .3, .1]))
    y = pm3.Categorical('y', p=tt.squeeze(arr_ys[a, b, :]))
1 Like

Thanks!