Shape Issue with custom logp in DensityDist

Small suggestion, I would perhaps replace that loop by:

    # Compute log ( p_s * Pr(y_i|s) )
    phi = pt.concatenate([p[None] for p in phi] , axis=0)  # PyTensor doesn't have `pt.vstack`
    phi = pt.log(phi) + pt.log(p[:, None])