Refactoring away from BaseTrace/MultiTrace

[michaelosthege]

Can I get a distribution of gut feelings around BaseTrace and MultiTrace?

👉 Are we in favor of refactoring the inner workings in sampling.py away from them and towards the McBackend interfaces?

👉 How important is backwards compat for passing pm.sample(trace=...)? (A hard break/switch from BaseTrace/MultiTrace to mcbackend.Backend is a lot easier to realize.)

Specifically, the tricky parts are the signatures of the various internal functions inside the sampling control flow as well as the different initialization behavior of BaseTrace vs mcbackend.Chain:

• A BaseTrace is initialized by a mix of its constructor logic and a call to .setup(). The constructor called once before (typically) the object is deep-copied for all chains. Then .setup() is called once per chain.
• With McBackend all metadata information is accumulated once for a Run. Via Run().init_chain(chain_number) the chains are then initialized all from the same metadata. No deepcopy involved.

[twiecki]

I'm not in favor of switching to McBackend, if anything we maybe should just switch to xarray internally.

[ricardoV94]

I would need more info to be able to discuss this. What are the pros, what are the drawbacks, and what other alternatives are out there. More importantly what immediate problem (if any) does this solve already.

[michaelosthege]

Pros ➕

Cleaner code in sampling.py

_choose_backend is called in many places in the control flow and the signatures and behaviors aren't very clear.
This is high-risk w.r.t. seeds and possibly a source of bugs like #5378.

With well-defined backend initialization flow (and eventual InferenceData conversion in mind) this should become a lot more maintainable.

Reducing costly redundancy

Every time a BaseTrace.__init__ is called, it compiles functions & recomputes the initial point just to evaluate shapes & types.
But in the end this is all identical for all chain, so it should be done exactly once.

Dynamic shape support

Determining variable shapes from an initial point in BaseTrace.__init__ isn't compatible with dynamic shapes at sampling time.
On the Aesara & aeppl side we're preparing for such scenarios already, but the PyMC storage can't do it.

xarray support (among others)

An XarrayBackend should be quite easy to implement, actually. Check out the NumPyBackend from mcbackend.backends and you'll find that all the needed metadata is already passed to the __init__ method.

Still, this is just one specific backend and it doesn't support dynamic shapes.
We should evaluate if there are situations where xarray gives better performance, but we shouldn't limit ourselves to one particular backend.

Sampler stats

Currently we don't pre-allocate memory for sampler stats. In McBackend the sampler stats are treated similar to variables.
For example, the mcbackend.NumPyBackend pre-allocates memory also for sampler-stats as long as there are no dynamic shapes involved.

Backend agnosticism

The code in sampling.py or even PyMC as a whole shouldn't need to know how draws & stats are stored.
NumPy? Dask? Xarray? MySQL? ClickHouse? HDF5? CSV?
Related: #5371

As long as the backend adheres to a well-defined interface it shouldn't concern the PyMC codebase.
In fact this is already the case with BaseTrace, but with the limitations I mentioned above.

McBackend really is just an evolution of BaseTrace into something that doesn't rely on deepcopy for initialization and is more compatible with backends and InferenceData conversion by design.

Cons ➖

MultiTrace and sampling.py are currently tightly coupled to BaseTrace.
A good compromise might be an adapter around mcbackend.Run to get a MultiTrace-like interface, but the initialization flow of BaseTrace and mcbackend.Run/Chain is so different that supporting both in pm.sample(trace=...) is really hard.

[twiecki]

Thanks, that's helpful. Would using the NumPyBackend be basically like MultiTrace?

[michaelosthege]

Yes, we'd have NumPyBackend be the default backend and users could override that either via kwarg, or even with something like pm.rcParams["backend"] = mcbackend.ClickHouseBackend("arviz-warehouse.mycompany.com")

[michaelosthege]

A mcbackend.Run currently doesn't have methods for fetching data (other than .to_inferencedata()), but we can definitely add some if we have good ideas which ones make sense.
The slicing API of MultiTrace doesn't have a very clear signature and gave us headaches in the past. Also it probably doesn't generalize well for dynamic shapes.
Access to draws/stats is currently via the individual Chains: https://github.com/michaelosthege/mcbackend/blob/v0.1.1/mcbackend/core.py#L73-L87

Feel free to propose method signatures!