Add find_MAP with close JAX integration and fix bug with Laplace fit

[jessegrabowski]

Closes #376

This PR adds code to run find_MAP using JAX. I'm using JAX for gradients, because I found the compile times were faster. Open to suggestions/rebuke.

It also adds a fit_laplace function, which is bad because we already have a fit_laplace function. This one has slightly different objective though -- it isn't meant to be used as a step sampler on a subset of model variables. Instead, it is meant to be used on the MAP result to give an approximation to the full posterior. My function also lets you do the Laplace approximation in the transformed space, then do sample-wise reverse transformation. I think this is legit, and lets you obtain approximate posteriors that respect the domain of the prior. Tagging @theorashid so we can resolve the differences.

Last point is that I added a dependency on better_optimize. This is a package I wrote that basically rips out the wrapper code used in PyMC find_MAP and applies it to arbitrary optimization problems. It is more feature complete than the PyMC wrapper -- it supports all optimizer modes for scipy.optimize.minimize and scipy.optimize.root, and also helps get keywords to the right place in those functions (who can ever remember if an argument goes in method_kwargs or in the funciton itself?). I plan to add support for basinhopping as well, which will be nice for really hairy minimizations.

I could see an objection to adding another dependency, but 1) it's a lightweight wrapper around functionality that doesn't really belong in PyMC anyway, and 2) it's a big value-add compared to working directly with the scipy.optimize functions, which have gnarly, inconsistent signatures.

[theorashid]

Hey, nice one, yeah I agree, we should only have one fit_laplace function.

it isn't meant to be used as a step sampler on a subset of model variables

The current fit_laplace isn't either. It isn't a step sampler. (The INLA stuff #340 still has a few blockers so that's separate and not yet in the library.) The implementation was made by a user following Statistical Rethinking, where McElreath fits some models using the Laplace approximation of all parameters.

Current behaviour of fit_laplace is:

• Take the MAP estimate of the parameters (https://github.com/pymc-devs/pymc-experimental/blob/main/pymc_experimental/inference/laplace.py#L109)
• Move to untransformed space (https://github.com/pymc-devs/pymc-experimental/blob/main/pymc_experimental/inference/laplace.py#L112)
• Get the Hessian (https://github.com/pymc-devs/pymc-experimental/blob/main/pymc_experimental/inference/laplace.py#L114)
• cov = np.linalg.inv(hessian) https://github.com/pymc-devs/pymc-experimental/blob/main/pymc_experimental/inference/laplace.py#L119
• simulate from this posterior approx rng.multivariate_normal(mean, cov, size=(1, draws)) and turn into idata. (Looking at this now actually, we should probably transform back into the constrained space? Otherwise I can see e.g. standard deviation variables with samples below zero.)

The behaviour when you only pass a subset of variables isn't really desirable in my opinion (see #345 (comment)), so we put a warning. So as you say:

Instead, it is meant to be used on the MAP result to give an approximation to the full posterior.

Agree, that's the best plan for fit_laplace.

Judging by your docs and a quick glance at your code, I think you're basically doing the same thing. The current implementation is few lines of code and a few docs, so I reckon

1. make sure you can pass the test case with your method, which is an example from BDA3 https://github.com/pymc-devs/pymc-experimental/blob/main/tests/test_laplace.py
2. throw any of the useful code and docs into your method

Then it should be safe to delete the existing code and we can go back to one fit_laplace.

I could see an objection to adding another dependency

I would love a generic optimiser in p u r e pytensor, but I can see looking at your code that there a lot of fancy extras that would take a large effort to write in pytensor. Still, if we want to go back to one of our efforts with a fixed point operator (pymc-devs/pytensor#978 and pymc-devs/pytensor#944), we could probably write find_MAP with that in some form, with fewer bells and whistles though.

Happy to look at your code and review properly later in the week if you'd like me to. Let me know. Otherwise, I'll leave to the core devs.

[ricardoV94]

Happy to look at your code and review properly later in the week if you'd like me to. Let me know. Otherwise, I'll leave to the core devs.

That would be appreciated

[ricardoV94]

Agree with what @theorashid said. This fit_laplace is going for the same goal as the previous one. Happy to replace it, if it's not married to JAX backend. Still fine to allow using JAX for the autodiff. What you're offering is very similar to nutpie gradient_backend kwarg, so we could use the same terminology

[ricardoV94]

No objections about your custom library wrapper

[jessegrabowski]

tagging @theorashid -- I couldn't pick you as a reviewer?

I did a major refactor of this. I broke the marriage to jax and generalized the find_MAP function. Files have been renamed to reflect this. I also merged the two laplace approaches. The biggest change is that I removed the ability to choose vars. I think the idea here was to be able to partially marginalize some variables in a model? But I think this would require a somewhat different approach.

[theorashid]

yea sorry I'm just a normal, but I'll give it a review. Will do it at some point in the next 2 weeks.

[ricardoV94]

Minor suggestions, PR looks amazing!

[ricardoV94]

@jessegrabowski can we close #376 with this PR?

Do you have a test that covers something like it?

[jessegrabowski]

Yes, I think this test and this test should cover that issue

[theorashid]

sweet, all done?

[jessegrabowski]

For now, though I'd still appreciate it if you could have a look and open issues on any bugs/shortcomings you find

[theorashid]

I managed to follow the code through and it looks good to me. Happy you got rid of the option to fit on a subset of variables, which didn't make sense to me anyway. If it passes the original test then it should be good. You can do something about the other comments if you want, but maybe not because we are e x p e r i m e n t a l

[theorashid]

For my understanding, what sort of scenarios/models would get a not PSD hessian. And is using closest PSD a good ideas?