PyMC3 posterior prediction example; how to incorporate a matrix of betas?

Thanks for your time and help! It is much appreciated in helping me get the right structure so that it may be used for all other problems I have.

I am running into issues using pm.Data as it is not liking my dimensions, but I am not sure why:

# generate Likert y; 1-7
obs_y = np.random.randint(7, size=(100, 1)) + 1

# generate treatment and control group
treatment = np.random.randint(3, size=(100,))
# convert to one-hot;
treatment = pd.get_dummies(data=treatment, drop_first=True)

# generate many categorical controls
covariates = np.random.randint(2, size=(100, 56))

# treatment to xarray
treatment = treatment.to_xarray().to_array()  # tried both; to_xarray()
treatment.shape  # (2, 100)

# coords for pymc3
coords = {
    'obs_id': np.arange(len(obs_y)),
    'treatment_betas': treatment  # (2, 100)
    }

# specify model
with pm.Model(coords=coords) as main_lm_model:
    t = pm.Data('treatment_betas', treatment, dims=('treatment_betas'))

# Length of `dims` must match the dimensions of the dataset. (actual 1 != expected 2)