I would try subsetting the beta
“matrix” by group (row), i.e. (beta[indexed_states] * x).sum(axis=1)
. The expression x.dot(beta[indexed_states]
might also work.
I was also struggling with this a while ago, cf. Unexpected difference in posterior estimates from two equivalent linear hierarchical models, fyi.
Unrelated: I found that pymc
can work directly with pandas.DataFrame
and Series
objects so no need to call .values
(if you want to convert use .to_numpy()
).