ENH: Propose new mutable_coords kwarg for pm.Model

[drbenvincent]

Before

import pymc as pm
import numpy as np

x = np.arange(20)
y = x + np.random.normal(scale=0.5, size=20)
xtrain, ytrain = x[:15], y[:15]
xtest, ytest = x[15:], y[15:]

# in sample parameter estimation
with pm.Model() as model:
    x = pm.MutableData('x', xtrain, dims='obs_ind', coords={"obs_ind": np.arange(len(xtrain))})
    y = pm.MutableData('y', ytrain, dims='obs_ind')
    alpha = pm.Normal('alpha', mu=0, sigma=10)
    beta = pm.Normal('beta', mu=0, sigma=10)
    sigma = pm.HalfNormal('sigma', 1)
    mu = alpha + beta * x
    pm.Normal('y_obs', mu=mu, sigma=sigma, observed=y, dims='obs_ind')
    idata = pm.sample()

# out of sample prediction  
with model:
    pm.set_data({"x": xtest, "y": ytest}, coords={"obs_ind": np.arange(len(xtest))})
    idata.extend(pm.sample_posterior_predictive(idata))

The thing which is perhaps non-obvious from a user perspective is providing coords via the pm.MutableData, which is totally different from when you provide coords via pm.Model.

After

# in sample parameter estimation
with pm.Model(mutable_coords={"obs_ind": np.arange(len(xtrain))}) as model:
    x = pm.MutableData('x', xtrain, dims='obs_ind')
    y = pm.MutableData('y', ytrain, dims='obs_ind')
    alpha = pm.Normal('alpha', mu=0, sigma=10)
    beta = pm.Normal('beta', mu=0, sigma=10)
    sigma = pm.HalfNormal('sigma', 1)
    mu = alpha + beta * x
    pm.Normal('y_obs', mu=mu, sigma=sigma, observed=y, dims='obs_ind')
    idata = pm.sample()
    
# out of sample prediction  
with model:
    pm.set_data({"x": xtest, "y": ytest}, coords={"obs_ind": np.arange(len(xtest))})
    idata.extend(pm.sample_posterior_predictive(idata))

Context for the issue:

Currently, the API for doing out of sample inference with coords (that change dimensionality) is not the easiest. You have to specify the coords as mutable by adding as a kwarg within the pm.MutableData, which is a bit non-obvious and potentially a little clunky.

I don't propose that we remove this, but the API which is really begging to exist is to be able to provide a mutable_coords kwarg to pm.Model.

[twiecki]

Did you see #6515?

[drbenvincent]

Whoops, I searched for "mutable" and "coords" but only for open issues.

So I think my suggestion is similar to the one @michaelosthege proposed here #6515 (comment)

Looking at the changes in the already merged #6515, I think this already implements my proposal. So I'll close this.