How to include coords in out-of-sample prediction?

Questions v5
2

Are you looking for something like this?


coords = {"coeffs": labels,
          "ydim":range(y_train.size)}

with pm.Model(coords=coords) as model:
    # data containers
    X = pm.MutableData("X", x_train)
    y = pm.MutableData("y", y_train)
    # priors
    b = pm.Normal("b", mu=0, sigma=1, dims="coeffs")
    # linear model
    mu = pm.math.dot(X, b)
    # link function
    p = pm.Deterministic("p", pm.math.invlogit(mu))
    # likelihood
    pm.Bernoulli("obs", p=p, observed=y, dims="ydim")


with model:
    idata = pm.sample()

with model:
    pm.set_data({"X": x_test, "y": y_test})
    model.set_dim("ydim", y_test.size, coord_values=range(y_test.size))
    predictions = pm.sample_posterior_predictive(trace=idata, predictions=True).predictions