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