Shape Issue with custom logp in DensityDist

Questions v5
5

Sure thing:

import arviz as az
import numpy as np
import pandas as pd
import pymc as pm
from pytensor import tensor as pt
from numpy.random import default_rng

RANDOM_SEED = 1234
np.random.seed(RANDOM_SEED)
rng = default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")
az.rcParams["stats.hdi_prob"] = 0.89

df = pd.read_csv("https://raw.githubusercontent.com/rmcelreath/rethinking/master/data/Boxes.csv", delimiter=";")

def logp(value, p, majority_first):

    phi = [None] * 5

    # Probability of data
    phi[0] = pt.switch(pt.eq(value, 2), 1, 0)
    phi[1] = pt.switch(pt.eq(value, 3), 1, 0)
    phi[2] = pt.switch(pt.eq(value, 1), 1, 0)
    phi[3] = pt.ones_like(value) * 1 / 3
    phi[4] = pt.switch(
        pt.eq(majority_first, 1), pt.switch(pt.eq(value, 2), 1, 0), pt.switch(pt.eq(value, 3), 1, 0)
    )

    # Compute log ( p_s * Pr(y_i|s) )
    for i in range(5):
        phi[i] = pt.log(phi[i]) + pt.log(p[i])

    # Compute average log-probability of y_i
    return pt.sum(pm.math.logsumexp(phi, axis=0))


with pm.Model() as m16_2:
    # prior
    p = pm.Dirichlet("p", np.array([4, 4, 4, 4, 4]))

    # majority first data
    majority_first = pm.ConstantData("majority_first", df["majority_first"].values)

    # likelihood
    y = pm.DensityDist(
        "y", p, majority_first, logp=logp, ndims_params=[1, 0], observed=df["y"].values
    )

    idata = pm.sample()