theano - 階層型pymc3モデルでラムダ関数を作成できません

Question

以下に示すモデルを PyMC 3 で作成しようとしていますが、ラムダ関数を使用して確率を観測データに適切にマッピングする方法がわかりません。

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)]

私が得るエラーは

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

私が構築しようとしているモデルでは、の要素は、zどの要素がphi対応する観測値の確率を与えるかを示しますdata( RV に配置されますw)。言い換えると、

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

Theano 式で確率を定義するか、Theano を使用する必要があるas_opと思いますが、このモデルでそれを行う方法がわかりません。