Error when using pm.Potential with pm.Data + coordinates

Hi everyone!

I am trying to use pm.Data and the coords keyword but I’m getting an error when sampling. Here’s the model:

coords = {
    'LoT': np.arange(len(L)),
    'cat': np.arange(L.shape[1]),
    'obs': np.arange(len(category_i))

with pm.Model(coords=coords) as model:
    L_data = pm.Data('L', L)#, dims=('LoT', 'cat'))
    cat_i_data = pm.Data('cat_i', category_i)#, dims='obs')
    out_i_data = pm.Data('outcome_i', outcome_i)#, dims='obs')
    L_dims = L.shape
    n_lots, n_cats = L.shape

    sigma = pm.HalfNormal('sigma', sigma=2, dims='cat')
    a_0 = pm.Normal('a_0', mu=0, sigma=2)
    a_1 = pm.Normal('a_1', mu=0, sigma=2)

    z_logp = tt.zeros((n_lots,), dtype='float')
    for z_i in range(n_lots):
        zi_logp = 0.
        for cat in range(n_cats):
            obs_idx = np.where(cat_i_data==cat)[0]
            muZ = a_0 + a_1 * L[z_i, cat]
            zi_logp = zi_logp + tt.sum(
                .dist(mu=muZ, sd=sigma[cat])
        z_logp = tt.set_subtensor(z_logp[z_i], zi_logp)
    lp3 = pm.Deterministic('z_logp', z_logp, dims='LoT')
    tot_logp = pm.math.logsumexp(z_logp)
    pot = pm.Potential('e', tot_logp)

The model compiles fine, but when I try to sample:

And in the trace zlog_p is always 0, so clearly the sampling isn’t working.

However, when I use the arrays directly instead of using their pm.Data version, everything works fine, and when I comment out

    lp3 = pm.Deterministic('z_logp', z_logp, dims='LoT')
    tot_logp = pm.math.logsumexp(z_logp)
    pot = pm.Potential('e', tot_logp)

everything also works fine (well, except that it’s a different model, but I mean that it runs fine). Also, when I comment out just pot = pm.Potential('e', tot_logp) the problem is still there. So my guess is that it has something to do with the Deterministic z_logp, but I’m not sure what about it is problematic. Thanks in advance for any ideas!

1 Like

I am not surprised the deterministic does not handle dims. It is not really supposed to do anything other than store the result of deterministic operations.


Is the issue solved by using dims in pm.Data but not in pm.Deterministic?

I thought dims in deterministics were only stored and passed downstream to ArviZ at conversion but had no effect within PyMC3

Unfortunately not! Even when changing the line to

    lp3 = pm.Deterministic('z_logp', z_logp)#, dims='LoT')

It gives the same error.

To be specific: if I use arrays directly instead of the pm.Data objects, everything works fine. So I don’t think the problem is about Deterministic and dims. It must have something to do with the combination of Deterministic, dims, and pm.Data.

(To generate fake data:)

sigma = 0.1
a_0, a_1 = 0, 1
z = 3 
L = np.array([
        [1, 2, 1],
        [3, 2, 1],
        [1, 4, 5],
        [5, 6, 8]

# category of the ith observation
category_i = np.repeat(np.arange(L.shape[1]), 1000)
mu_i = a_0 + a_1 * L[z,category_i]
outcome_i = np.random.normal(