It looks like you are using an array with one dimension of length obs_id so you effectively can’t use obs_id, treatment as dims, your array only has one dimension. You should be using an array with two dimensions obs_id, treatment and have the treatment dimension have length 1. An easy way to do that from a 1d treatment array is treatment[:, None]
As for the error, I don’t really understand it, treatment_beta looks fine to me and I don’t know what i_1_0 is.