Refactor `init` logic in `pm.sample`

[jessegrabowski]

Description

We currently have a long list of if/else conditions. but it's needlessly restrictive. Basically the user is allowed to select an initialization mode (jitter, MAP, or advi) and an adaptation mode (adapt_diag, adapt_diag_grad, or adapt_full). Currently some of these combinations are arbitrarily disallowed, for example advi+adapt_diag_grad seems nice.

I propose splitting the incoming string on + and then dispatching to appropriate functions. This would make it easier to implement new initialization or adaptation methods. I can imagine the nutpie fisher loss adaptation, or normalizing flow adaption (or initialization!), being added as choices, for examples.

[Chirag3841]

@jessegrabowski I guess current if/else chain is unnecessarily restrictive and hard to extend. We can split the string on "+", identify init mode (jitter/map/advi) and adaptation mode (adapt_diag/adapt_diag_grad/adapt_full) independently, then dispatch via small registries/dicts. This would naturally allow combinations like advi+adapt_diag_grad and make adding future modes much easier.

[velochy]

Logic definitely makes sense.
I can try taking a stab at this along w adding pathfinder as an init option (once the refactor at pymc-devs/pymc-extras#656 is complete and merged)

[jessegrabowski]

Definitely in favor of allowing pathfinder, but I'd like to find a general solution to the problem here.

Right now, init does 3 things:

1. It lets you choose an initialization function, which is a map from a logp graph to a mean and covariance matrix
2. It helps you add random jittier to the mean returned by (1), if you want
3. It lets you choose an adaptation strategy for NUTS

These are 3 very different things. 1 and 2 aren't even NUTS specific, but init only works with NUTS. We can provide some strings (right now we provide MAP, advi, etc), but it is not possible to do configuration of the chosen method (what if you want a different optimizer in ADVI? A different minimize algorithm in MAP? etc). We can provide some pre-packaged options with smart defaults as strings (which can include pathfinder), but I think what I want from this issue is a more general solution.

[jessegrabowski]

@jessegrabowski I guess current if/else chain is unnecessarily restrictive and hard to extend. We can split the string on "+", identify init mode (jitter/map/advi) and adaptation mode (adapt_diag/adapt_diag_grad/adapt_full) independently, then dispatch via small registries/dicts. This would naturally allow combinations like advi+adapt_diag_grad and make adding future modes much easier.

Agreed we should have some kind of registry. The values in the registry should be some kind of API that users can also pass directly if they want.

[Chirag3841]

@jessegrabowski I would like to solve this issue if anyone had not taken it yet too.

[velochy]

1. It lets you choose an initialization function, which is a map from a logp graph to a mean and covariance matrix

I have a slight issue with this. Namely, Im working with models that often compile for 10-20min. Passing graph here means it gets compiled twice: once inside the init, second time fir sample.

Why not pass in compiled logp, dlogp?

[jessegrabowski]

Why not pass in compiled logp, dlogp?

No objection to that. My thinking was that we want all of the difference inference methods in pymc-extras to just work out of the box as valid initializers.

Ideally we'd work hard to get those compile times down. But that's obviously neither here nor there for this PR.

[velochy]

@jessegrabowski I would like to solve this issue if anyone had not taken it yet too.

@Chirag3841 feel free to take the lead, as it would realistically be a week or two before I could get around to it.

As noted above, my main interest is getting pathfinder to work as an init method, so I have some opinions on what the interface on initializer side might look like, so I'll be following the work you do closely and rooting for you :)

[velochy]

Why not pass in compiled logp, dlogp?

No objection to that. My thinking was that we want all of the difference inference methods in pymc-extras to just work out of the box as valid initializers.

Ideally we'd work hard to get those compile times down. But that's obviously neither here nor there for this PR.

Have you considered a compile cache? Hashing a pytensor graph might be tricky but this is something that might cut some boilerplate out of multiprocessing use cases so maybe something to think about?

[jessegrabowski]

Yes i've thought about something like this. Or at least a rewrite cache. Sometimes when you have a model that compiles really slowly it's because it's doing tree(3) fusion rewrites. This is especially true if you have lots of scalar inputs / stacked equations. If we had an LRU cache on the rewriter we could ease some of that. I've even thought we should consider a rust-based rewriter, but the downside there is that it becomes a lot harder to add new rewrites. Being pure python is a pretty big upside in terms of hackability/accessibility.

I'm sure @ricardoV94 would have a lot more and more developed thoughts on this.