Update and Refactor find_MAP and fit_laplace

[jessegrabowski]

I just made an update to better-optimize that uses hessian matrix caching for better performance. This is something we can immediately take advantage of with find_MAP, re-using sub-computation from the loss or gradient in the hessian computation. This PR updates the functions generated by find_MAP to take advantage of this.

While I was at it, I went ahead and did some cleanup and reorganization of the code. In particular:

1. I split up the files into smaller, more logical groupings. I moved all the files to a laplace_approx submodule.
2. find_MAP now returns an idata. This is more consistent with all the other PyMC sampling function -- it's weird to get back a dictionary in this one case.
3. When available, find_MAP will now always store the inverse hessian. This is done to try to avoid an extra function compilation when it is used in conjunction with fit_laplace.
4. fit_laplace was a really dumb function that was inexplicably sampling from scipy distributions. This required a ton of unnecessary work. If only we had a PPL that could help sample from complicated distributions...

fit_laplace still isn't perfect. I wanted to store both the value variables and the transformed RVs as deterministics in a pymc model and sample them directly, but that doesn't appear to work -- maybe this is a bug? I ended up doing two passes, once for the constrained RVs, then a second pass for the unconstrained. It would be good to minimize that.

I also removed as many little options that were floating around as possible. These function signatures were already horrible.

Finally, I eliminated a lot of test parameterizations to speed the CI up, but also added a lot of new tests for functions that were previously not covered. Hopefully it's still net positive.

[jessegrabowski]

One slightly janky thing is that I store the MAP result in the posterior group, but without chain and draw dimensions. That might be a bad choice, because it breaks a promise that idata typically makes. But they also don't have these dimension. Not sure.

This is no longer true, find_MAP returns dummy chain, draw dims now. I thought it was too much to break the arviz promise that posterior always has chain/draw

Another jank choice is the temp_chain, temp_draw thing in model_to_laplace_approx. Someone else might know a better way to accomplish what I'm doing here :)

[Copilot]

Pull Request Overview

This PR refactors the MAP-finding and Laplace approximation routines to improve performance by caching Hessian computations, reorganizes code into a laplace_approx submodule, and standardizes return types to ArviZ InferenceData.

• Cache and reuse Hessian subcomputations in find_MAP/fit_laplace workflows.
• Move all Laplace-related modules under pymc_extras/inference/laplace_approx.
• Update find_MAP to return InferenceData and simplify fit_laplace interface.

Reviewed Changes

Copilot reviewed 15 out of 18 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/inference/laplace_approx/test_find_map.py	Add tests for new find_MAP defaults, PSD helper, and JAX paths
pymc_extras/inference/laplace_approx/scipy_interface.py	New module for compiling loss/grad/Hessian for SciPy optimizers
pymc_extras/inference/laplace_approx/laplace.py	Refactor fit_laplace, cache inverse Hessian, build idata
pymc_extras/inference/laplace_approx/find_map.py	Refactor find_MAP, wrap into InferenceData, split logic
pymc_extras/inference/laplace_approx/idata.py	Helpers to add data/fit/optimizer results into InferenceData
pymc_extras/inference/pathfinder/pathfinder.py	Update import to new add_data_to_inference_data helper
pymc_extras/inference/fit.py	Route fit(method="laplace") to new Laplace submodule
pyproject.toml	Bump better-optimize dependency to ≥0.1.4

Comments suppressed due to low confidence (2)

pymc_extras/inference/laplace_approx/scipy_interface.py:101

• The docstring lists f_fused and f_hessp as return values but the function actually returns a list of one or two Function objects. Update the doc to reflect that it returns a list[Function] (or [Function, Function]).

    f_fused: Function

pymc_extras/inference/laplace_approx/scipy_interface.py:53

• The return statement uses a starred expression (return *loss_and_grad, hess), which is invalid syntax in Python. Wrap the unpacking in a tuple, e.g.: return (*loss_and_grad, hess).

            return *loss_and_grad, hess

[jessegrabowski]

This should be ready to go. Last changes:

• find_laplace should work in all cases, including when the value variables have a different shape as the RVs (added test for this). I brought back the old way of making the batched RVs, then combine it with the new pymc model approach to make the raveled vector. I think the result is quite nice.

• find_MAP posterior has chain and draw, to stick fast to the arviz convention (but it can be easily squeezed away by the user)

• I found some unused helpers and removed them. Also removed the utilities file, because I'm sick of those and I never remember where I put anything.

[ricardoV94]

arviz suggests a separate group for these cases https://python.arviz.org/en/latest/schema/schema.html#unconstrained-posterior

[ricardoV94]

But we don't have coords for those ofc

[jessegrabowski]

That would be for all the unconstrained values yeah, not just the oddball ones?

[ricardoV94]

No idea. I don't see why we have this function though (see my other comment)

[ricardoV94]

Why do we have this function? Didn't you make a model where each variable is already an unstacked deterministic?

[jessegrabowski]

Let me double-check, might be a holdover from the old stuff that I got confused about.

[ricardoV94]

The unobserved_value_vars graph already converts from latent space to constrained space, then you vectorize that with batch draws and you have everything. That's how I read it. Which if true, is nice, this function call all go, and you don't worry about coords, since you don't store constrained draws per RV, only the whole concatenated vector?

[jessegrabowski]

Ok looking it back over, the purpose of this function is to take the draws from the actual laplace approximation and also return those as part of the posterior. My logic is that users might want this for diagnostic purposes, since this is where the multivariate normal actually lives.

My impression (based on very little) is that most packages won't do the constraining transformation on the outputs at all. I have this impression because people often cite "not respecting the domain of the priors" as a reason why laplace isn't the best tool.

[ricardoV94]

Isn't that just the long flatten vector?

[ricardoV94]

Okay. Still seems a niche place to invent this stuff here

[jessegrabowski]

I also agree with that.

I was thinking that the right solution for this is to add a method to transformers that acts on coords. Conceptually, that seems like the right place for this.

[jessegrabowski]

On the strategy point you don't need two steps, when you did the flat->constrained mapping using unobserved_value_names as outputs you could have done
flat->constrained+unconstrained mapping, using value_vars + unobserved_value_vars as outputs.

Where is this? In the call to join_nonshared_inputs?

[ricardoV94]

yeah

[jessegrabowski]

pm.sample_posterior_predictive filters out value variable names here, so doing everything symbolically in one shot doesn't work.

I'm sticking with what I have for now, I really just need this PR to be done.