MLDA needs refactoring to become operable again

[michaelosthege]

Description of your problem

In #5019 we removed the trace continuation feature, because it's making the logic in sampling.py considerably harder, particularly w.r.t. initial values.

The MLDA can probably be fixed easily by using pm.sample(start=...) and collecting the individual traces to concatenate them at the end.

@gmingas please check if you want to fix it.

Versions and main components

• PyMC3 Version: main

[gmingas]

Hi, thanks for pointing this out, sounds straightforward. @mikkelbue is maintaining this at the moment.

[mikkelbue]

Hi
Thanks for the mention. I made use of trace continuation in the first place because it sped up sampling considerably. I think it was because it avoided creating and destroying new traces all the time, but I am not sure. Anyway, it is non-essential for the functionality.

@gmingas Would it break variance reduction if we just go back to the way you had implemented it in the first place, where it doesn't collect the full coarse chain?

I will have a look when I find some time, but I am trying to wrap up my PhD at the moment, so things are a bit tight.

[mikkelbue]

Hi all. Just to let you know, we are working on this, but after solving this problem, a new issue has come up with the variance reduction feature, see alan-turing-institute#95.

[mikkelbue]

Hi @michaelosthege
Just a short comment after running some tests. The trace continuation is not really the most important feature of the subsample function. They key is that it avoids a lot of the overhead of the pm.sample function. Essentially, it is a thin wrapper around _iter_sample() that skips all the tests and multi-chain tricks of pm.sample and _sample_many.

I tried using pm.sample, but using my refactored version of subsample is about 10x faster. So I would like to stick with the boiled-down version, but of course avoiding trace continuation. Unless there are some other breaking changes coming up about the architecture of the pm.sample and _iter_sample.

[michaelosthege]

There are breaking charges on the lower-level functions, specifically around the start/initval kwarg. Overall the functions now require the dict to be complete and no longer run checks, so it should actually become a little faster. See #4983 and #5027

[mikkelbue]

Thank you for your comments. Sorry for taking so long. I was waiting for the big PR you mentioned to be merged to be sure that this would fit in.

As you can see I have created a pull request with the required changes, and reintroduced the MLDA tests in test_step.py.