Unable to create lambda function in hierarchical pymc3 model

Question

I'm trying to create the model shown below with PyMC 3 but can't figure out how to properly map probabilities to the observed data with a lambda function.

import numpy as np
import pymc as pm

data = np.array([[0, 0, 1, 1, 2],
                 [0, 1, 2, 2, 2],
                 [2, 2, 1, 1, 0],
                 [1, 1, 2, 0, 1]])
(D, W) = data.shape
V = len(set(data.ravel()))
T = 3
a = np.ones(T)
b = np.ones(V)

with pm.Model() as model:
    theta = [pm.Dirichlet('theta_%s' % i, a, shape=T) for i in range(D)]
    z = [pm.Categorical('z_%i' % i, theta[i], shape=W) for i in range(D)]
    phi = [pm.Dirichlet('phi_%i' % i, b, shape=V) for i in range(T)]
    w = [pm.Categorical('w_%i_%i' % (i, j),
                        p=lambda z=z[i][j], phi_=phi: phi_[z], # Error is here
                        observed=data[i, j])
            for i in range(D) for j in range(W)]

The error I get is

AttributeError: 'function' object has no attribute 'shape'

In the model I'm attempting to build, the elements of z indicate which element in phi gives the probability of the corresponding observed value in data (placed in RV w). In other words,

P(data[i,j]) <- phi[z[i,j]][data[i,j]]

I'm guessing I need to define the probability with a Theano expression or use Theano as_op but I don't see how it can be done for this model.

Answer 1

You should specify your categorical p values as Deterministic objects before passing them on to w. Otherwise, the as_op implementation would look something like this:

@theano.compile.ops.as_op(itypes=[t.lscalar, t.dscalar, t.dscalar],otypes=[t.dvector])
def p(z=z, phi=phi):
    return [phi[z[i,j]] for i in range(D) for j in range(W)]