Out-of-Sample Predictions example is outdated for PyMC 6 / pymc-bart 0.12

[fingoldo]

Hi!

I think the "Out-of-Sample Predictions" section of the BART documentation is no longer correct for current releases.

Documentation:

https://www.pymc.io/projects/bart/en/latest/examples/bart_introduction.html#Out-of-Sample-Predictions

Environment

PyMC      6.1.0
pymc-bart 0.12.0

Problem

The documentation currently suggests:

with model:
    pm.set_data({"X": X_test})

    predictions = pm.sample_posterior_predictive(
        idata,
        var_names=["mu"],
        predictions=True,
    )

However, this produces predictions with the training-set length, not the test-set length.

Interestingly, the model itself is updated correctly:

print(model["X"].eval().shape)
print(model["mu"].eval().shape)

outputs

(20, 2)
(20,)

so pm.set_data() works correctly.

The problem appears to happen inside sample_posterior_predictive(), which returns

(2, 50, 100)

instead of

(2, 50, 20)

Minimal reproducer

import numpy as np
import pymc as pm
import pymc_bart as pmb

print("PyMC:", pm.__version__)
print("pymc_bart:", pmb.__version__)

rng = np.random.default_rng(42)

X_train = rng.normal(size=(100, 2)).astype(np.float32)
y_train = (
    np.sin(X_train[:, 0])
    + X_train[:, 1] ** 2
    + rng.normal(scale=0.1, size=100)
).astype(np.float32)

X_test = rng.normal(size=(20, 2)).astype(np.float32)

with pm.Model() as model:

    X = pm.Data("X", X_train)

    mu = pmb.BART("mu", X, y_train)

    sigma = pm.HalfNormal("sigma", 1)

    pm.Normal(
        "y",
        mu=mu,
        sigma=sigma,
        observed=y_train,
    )

    trace = pm.sample(
        draws=50,
        tune=50,
        chains=2,
        random_seed=42,
        progressbar=False,
        return_inferencedata=True,
    )

with model:
    pm.set_data({"X": X_test})

    print("X shape :", model["X"].eval().shape)
    print("mu shape:", model["mu"].eval().shape)

    pred = pm.sample_posterior_predictive(
        trace,
        var_names=["mu"],
        predictions=True,
        progressbar=False,
        return_inferencedata=False,
    )

print(pred["mu"].shape)

Output:

PyMC: 6.1.0
pymc_bart: 0.12.0

X shape : (20, 2)
mu shape: (20,)

(2, 50, 100)

Expected output:

(2, 50, 20)

Workaround

The example starts working correctly after replacing

pm.sample_posterior_predictive(
    trace,
    var_names=["mu"],
    predictions=True,
)

with

pm.sample_posterior_predictive(
    trace,
    sample_vars=["mu"],
    predictions=True,
)

which returns

(2, 50, 20)

as expected.

Possible explanation

From the current PyMC documentation, it appears that sample_posterior_predictive() now reuses variables already stored in the posterior trace unless they are explicitly listed in sample_vars.

If that's the intended behavior, then the BART documentation example predates this API change and should probably be updated accordingly.

Would it make sense to update the "Out-of-Sample Predictions" section to use:

pm.sample_posterior_predictive(
    trace,
    sample_vars=["mu"],
    predictions=True,
)

instead of the current example? Or, mb even better, mentioning both options depending on version. I spent quite a time getting my snippet to work.

[aloctavodia]

Thanks for reporting this! Would you like to submit a PR updating the example? The examples live in https://github.com/pymc-devs/pymc-examples