`Manager` lifetime bug causes `TypeError` in `sample_posterior_predictive` for out-of-sample predictions

[juanitorduz]

I had to work with Claude to find this nasty one. This is Claudes/ suggested solution. I made it work in my notebook by doing:

with model:
    idata = pm.sample(
        tune=1_000,
        draws=800,
        chains=6,
        cores=6,
        target_accept=0.8,
        random_seed=rng,
    )
    # pymc_bart stores trees in a multiprocessing Manager().list() (ListProxy).
    # The Manager is a local in BART.__new__ and gets GC'd, making the ListProxy
    # inaccessible for later OOS calls. Converting to a plain list here (while the
    # manager is still alive) fixes the TypeError in sample_posterior_predictive.
    type(mu.owner.op).all_trees = list(type(mu.owner.op).all_trees)
    idata.extend(pm.sample_posterior_predictive(trace=idata, random_seed=rng))

Summary

When calling pm.sample_posterior_predictive on new data (after pm.set_data) in a model that contains a BART variable, a cryptic TypeError is raised:

TypeError: only integer scalar arrays can be converted to a scalar index

The root cause is a resource-lifetime bug: the multiprocessing.Manager whose server process backs BARTRV.all_trees (a ListProxy) is never persisted, so it can be garbage-collected at any time after BART.__new__ returns. Once the server is dead, the ListProxy is a dangling reference and any access to it (including the not cls.all_trees boolean check in rng_fn) fails.

Environment

• pymc-bart: 0.11.0
• pymc: ≥ 5.28.0
• Python: 3.13
• OS: macOS (darwin)

Minimal reproducible example

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

rng = np.random.default_rng(42)

X_train = np.random.randn(100, 3)
y_train = np.random.randn(100)
X_test  = np.random.randn(20, 3)

with pm.Model() as model:
    x_data = pm.Data("x_data", X_train)
    mu = pmb.BART("mu", X=x_data, Y=y_train, m=50)
    sigma = pm.HalfNormal("sigma", sigma=1)
    obs = pm.Normal("obs", mu=mu, sigma=sigma, observed=y_train)

with model:
    idata = pm.sample(tune=100, draws=100, chains=2, cores=2, random_seed=rng)
    idata.extend(pm.sample_posterior_predictive(idata, random_seed=rng))  # works

# --- later in the script / notebook ---

with model:
    pm.set_data({"x_data": X_test})
    idata_oos = pm.sample_posterior_predictive(idata, random_seed=rng)  # FAILS

Error output (abbreviated):

File .../pymc_bart/bart.py:55, in BARTRV.rng_fn(cls, rng, X, Y, m, alpha, beta, size)
     52 if not size:
     53     size = None
---> 55 if not hasattr(cls, "all_trees") or not cls.all_trees:

TypeError: only integer scalar arrays can be converted to a scalar index

Note that the in-sample sample_posterior_predictive (called immediately after pm.sample()) works. The bug only surfaces when sample_posterior_predictive is called later (e.g. for a test set).

Root cause

In BART.__new__ (bart.py, lines 143–168):

# bart.py — BART.__new__
manager = Manager()                     # (1) starts a managed server process
instance_all_trees = manager.list()     # (2) ListProxy backed by that server

bart_op = type(
    f"BART_{name}",
    (BARTRV,),
    {
        "all_trees": instance_all_trees, # (3) stored as a class attribute
        ...
    },
)()
# (4) manager goes out of scope here — eligible for GC immediately

The manager object is never stored anywhere persistent (not on bart_op, not on the model). In CPython the reference count drops to zero when BART.__new__ returns, making manager eligible for garbage collection. When the GC eventually runs (or the reference is released), it calls Manager.__del__ → Manager.shutdown(), which terminates the server process.

The ListProxy stored in cls.all_trees is now a dangling proxy: it still exists as a Python object but any communication with the dead server fails.

Why in-sample PPX works: Python's GC is not instantaneous. The manager is collected lazily, so the server is still alive when pm.sample_posterior_predictive is called immediately after pm.sample(). By the time OOS prediction runs (later in the script, possibly after other allocations trigger GC), the server is gone.

Why the error message is confusing: The TypeError: only integer scalar arrays can be converted to a scalar index originates from the failed proxy communication inside not cls.all_trees, rather than the expected EOFError / BrokenPipeError, because of how the ListProxy's internal error handling interacts with the numpy machinery.

Proposed fixes

Fix 1 — Keep manager alive on the class (minimal change, preserves current design)

# bart.py — BART.__new__
bart_op = type(
    f"BART_{name}",
    (BARTRV,),
    {
        "all_trees": instance_all_trees,
        "_manager": manager,             # <-- add this line to keep manager alive
        ...
    },
)()

Storing manager as a class attribute ensures it stays alive for as long as the bart_op class exists (i.e. for the lifetime of the model).

Fix 2 — Replace Manager().list() with a plain list (preferred)

multiprocessing.Manager is only necessary for cross-process shared state. Looking at pgbart.py:

# pgbart.py — PGBART.astep
if not self.tune:
    self.bart.all_trees.append(self.all_trees)  # self.all_trees is a numpy array

The PGBART step method runs in the same process as the main thread in PyMC's default sampler. A regular list is sufficient:

# bart.py — BART.__new__  (remove Manager entirely)
instance_all_trees = []   # plain list — no Manager needed

bart_op = type(
    f"BART_{name}",
    (BARTRV,),
    {
        "all_trees": instance_all_trees,
        ...
    },
)()

This also eliminates the not cls.all_trees ambiguity for numpy arrays (plain list.__len__ is unambiguous) and removes a subprocess overhead.

Note: If BART sampling is ever moved to a true multiprocessing context where workers run in separate processes, the Manager would need to be reintroduced — but it would need to be stored persistently as described in Fix 1.

Workaround (user-side)

Convert the ListProxy to a plain list immediately after pm.sample(), while the Manager server is still alive:

with model:
    idata = pm.sample(...)
    # Materialise the ListProxy into a plain list before the Manager is GC'd
    type(mu.owner.op).all_trees = list(type(mu.owner.op).all_trees)
    idata.extend(pm.sample_posterior_predictive(idata, random_seed=rng))

# OOS prediction now works
with model:
    pm.set_data({"x_data": X_test})
    idata_oos = pm.sample_posterior_predictive(idata, random_seed=rng)

Additional notes

• The bug is not specific to out-of-sample prediction: any workflow that calls sample_posterior_predictive a significant time after model construction (e.g. loading a saved idata and predicting in a new session) will hit it, since all_trees is not persisted to the InferenceData object.
• This is separate from (but related to) the known limitation that BART trees are not serialised into idata, making true posterior-predictive replay from a saved trace impossible without the live model object.

[juanitorduz]

@aloctavodia have you encounter this before? Any opinion on the fix? I can create a PR.

[aloctavodia]

I have not. A PR is welcome! I'm not sure suggestion 2 will work, though.

[juanitorduz]

@aloctavodia I think we can close this one as we are now relying on the Rust version, right?