Pathfinder

[miclegr]

ref #157

[junpenglao]

Leaving the first round of comments, largely style wise. Need a bit more indepth review of the pathfinder implementation (vi/pathfinder.py)
There are a lot of places single letter variable names are used. Suggest changing to more informative name, for example x -> state, x0 -> initial_state, J -> param_dims, g -> grad

[rlouf]

@miclegr Thank you for contributing, this looks like it will be a valuable contribution to the library. Could you please rebase main onto your branch? When you checked out to create your branch the CI tests were not properly configured and thus are not running.

[miclegr]

@miclegr Thank you for contributing, this looks like it will be a valuable contribution to the library. Could you please rebase main onto your branch? When you checked out to create your branch the CI tests were not properly configured and thus are not running.

sure, did the rebasing now. Please note that when it comes to tests there's the problem I've outlined in my comment in #157 :

Unfortunately the L-BFGS optimizer is not working well when optimizing model's (negative) log-lilekihood while working in jax's default float32 mode, usually convergence fails. I've noticed that by simply dividing model's likelihood by number of observations (hence optimizing model's average likelihood) optimization converges. Unfortunately given blackjax design it's quite unnatural to ask for average likelihood. So in the end it's recommended to turn on double precision mode

A consequence of this is that, for running pathfinder tests, double precision mode is needed. Since double precision mode needs to be set at jax initialization time (see here), test suite should support some test in float32 mode and some in float64 mode (e.g. by running them in separate processes). It's feasible (see here for example) but not implemented at the moment. Right know i've just set the float64 mode for that tests: https://github.com/miclegr/blackjax/blob/128ce00bd2b28e06f79d126c9d3c097b6378ccc8/tests/test_pathfinder.py#L2-L3 But this wont work when running the full suite of tests. Happy to spend some time to get the multi float mode tests in place, if that's the solution.

[miclegr]

Tentative feature request to have L-BFGS optimization path exposed directly in jax, without the code duplication of this pull request: jax-ml/jax#10243

[junpenglao]

Per discussion in jax-ml/jax#10243, seems it is a better practice to use JAXopt - could you update the code and add a JAXopt dependency?

[miclegr]

Yep, I'm going to start working on that. Do you prefer JAXopt as a full requirement (i.e. as an entry in requirements.txt) or as an additional one (i.e. import within a try/except)?

[junpenglao]

let's make it a full requirement - I imagine we will use it in other optimization related inference method

[rlouf]

@miclegr is there anything we can do to help?

[miclegr]

Hi! I was finally able to complete the refactoring to make the code use jaxopt instead of jax.scipy.optimize.
I've also taken some time to make it work as much as possible with float32, by making ad ad-hoc adjustment to the stepsize of the LBFGS during training: https://github.com/miclegr/blackjax/blob/pathfinder/blackjax/vi/pathfinder.py#L290-L295

Now the tests run fine without the need to turn on the float64 mode, also the notebook :) So it should be mergeable.
Happy to do any remaining work

Michele

[codecov]

Codecov Report

Merging #194 (33ac071) into main (1393af3) will decrease coverage by 3.28%.
The diff coverage is 72.36%.

@@            Coverage Diff             @@
##             main     #194      +/-   ##
==========================================
- Coverage   98.40%   95.12%   -3.29%     
==========================================
  Files          36       39       +3     
  Lines        1383     1580     +197     
==========================================
+ Hits         1361     1503     +142     
- Misses         22       77      +55     

Impacted Files	Coverage Δ
blackjax/kernels.py	84.70% <29.72%> (-15.30%)	⬇️
blackjax/adaptation/pathfinder_adaptation.py	37.14% <37.14%> (ø)
blackjax/vi/pathfinder.py	94.30% <94.30%> (ø)
blackjax/adaptation/__init__.py	100.00% <100.00%> (ø)
blackjax/vi/__init__.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 1393af3...33ac071. Read the comment docs.

[junpenglao]

Great work @miclegr! I made some minor clean up - will merge once test pass.

[miclegr]

Thanks! and thanks for all the great work at pymc and tensorflow_probability 😉