Add wrapper for running blackjax pathfinder #72

twiecki · 2022-09-07T13:16:36Z

Lot's of help from @ricardoV94

CC @rlouf

pymc_experimental/tests/test_pathfinder.py

pymc_experimental/inference/pathfinder.py

pymc_experimental/tests/test_pathfinder.py

twiecki · 2022-09-08T09:04:45Z

Adressed comments.

pymc_experimental/inference/pathfinder.py

pymc_experimental/tests/test_pathfinder.py

twiecki · 2022-09-08T13:40:36Z

Ready for next round.

ricardoV94

Looks good, just two hard-coded values still in there.

pymc_experimental/inference/pathfinder.py

ricardoV94 · 2022-09-08T15:42:10Z

Just checking if the tests pass, they were failing weirdly before.

ricardoV94 · 2022-09-08T15:42:56Z

Maybe you just need to rebase from main, because @ferrine did some changes? (if they fail again)

…function.

Co-authored-by: Ricardo Vieira <[email protected]>

pymc_experimental/bart/__init__.py

twiecki · 2022-09-08T16:14:01Z

Oh, of course it doesn't work on windows...

ColCarroll · 2022-09-08T16:23:18Z

Not sure who this comment is for, but:

it'd be cool to be able to use the covariance estimates from the lbfgs iterations to initialize a mass matrix for NUTS/HMC
it'd be even cooler to use this as an automatic initialization for those samplers: you could use the condition number of the covariance, and the desired acceptance probability to also estimate an initial step size. i'm specifically suggesting having a path towards providing pm.sample(init='pathfinder'). while in the experimental directory, it perhaps makes most sense to have a sample_pathfinder function?

neither of these are blockers, just hopeful suggestions!

ColCarroll · 2022-09-08T16:40:48Z

Sorry, looking more at the blackjax API, it seems like you'd get both the covariance and the initial values from just init, so I'm describing a different function from this one:

this one does a variational fit + generates (approximate) samples
i'd like to do a more efficient job initializing MCMC and generating (asymptotically, but not actually) exact samples

rlouf · 2022-09-08T17:58:45Z

it'd be cool to be able to use the covariance estimates from the lbfgs iterations to initialize a mass matrix for NUTS/HMC

We have an adaptation routine that does just that. We really need documentation 🙃

ricardoV94 · 2022-09-08T18:12:51Z

Oh, of course it doesn't work on windows...

pytest.skip based on sys.platform.

ricardoV94 · 2022-09-09T11:33:57Z

Goddamn, isn't there a way to maybe turn the import into a noop?

Good old try/except

twiecki · 2022-09-09T13:43:54Z

It's green :party:!

ricardoV94

What about this instead?

pymc_experimental/inference/fit.py

Co-authored-by: Ricardo Vieira <[email protected]>

ricardoV94 · 2022-09-09T14:28:21Z

pymc_experimental/inference/fit.py

+        try:
+            from pymc_experimental.inference import fit_pathfinder
+        except ImportError as exc:
+            raise RuntimeError("Need JAX/ Blackjax / wahever to use `pathfinder`") from exc


This is not acceptable

ricardoV94 · 2022-09-10T08:55:07Z

It's green :party:!

Haha it was green because no tests were run 🤦

twiecki · 2022-09-10T09:04:35Z

FML.

This reverts commit 75aa37b.

ghost reviewed Sep 7, 2022

View reviewed changes

pymc_experimental/tests/test_pathfinder.py Outdated Show resolved Hide resolved

ghost reviewed Sep 7, 2022

View reviewed changes

pymc_experimental/inference/pathfinder.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Sep 7, 2022

View reviewed changes

pymc_experimental/inference/pathfinder.py Outdated Show resolved Hide resolved

pymc_experimental/inference/pathfinder.py Outdated Show resolved Hide resolved

pymc_experimental/tests/test_pathfinder.py Outdated Show resolved Hide resolved

ricardoV94 requested changes Sep 8, 2022

View reviewed changes

pymc_experimental/inference/pathfinder.py Outdated Show resolved Hide resolved

pymc_experimental/tests/test_pathfinder.py Show resolved Hide resolved

ricardoV94 requested changes Sep 8, 2022

View reviewed changes

pymc_experimental/inference/pathfinder.py Outdated Show resolved Hide resolved

pymc_experimental/inference/pathfinder.py Outdated Show resolved Hide resolved

twiecki and others added 11 commits September 8, 2022 17:47

Add wrapper for running blackjax pathfinder.

26d4db0

Run black.

9cc95ec

Run precommit.

1e8ad4a

Add blackjax to requirements.

a4cf339

Do not make import optional.

43b6f8e

Add more kwargs. Add license. Improve tests. Add doc string. Add fit …

3a3b2d7

…function.

Add fit function to base namespace.

74de0f9

Update copyright year.

d4e9ab4

Add type to random_seed and better init. Test for correct shapes.

1bf3473

Update pymc_experimental/inference/pathfinder.py

79e89df

Co-authored-by: Ricardo Vieira <[email protected]>

Update pymc_experimental/inference/pathfinder.py

03406cc

Co-authored-by: Ricardo Vieira <[email protected]>

twiecki force-pushed the blackjax_pathfinder branch from afc09ea to 03406cc Compare September 8, 2022 15:47

ghost reviewed Sep 8, 2022

View reviewed changes

pymc_experimental/bart/__init__.py Outdated Show resolved Hide resolved

twiecki added 2 commits September 8, 2022 18:00

Fix import of fit function.

4f0dc4e

Add fit.py.

2cfeed9

Skip on windows.

746d1d1

twiecki added 2 commits September 9, 2022 15:07

try/except blackjax import.

b6d2843

try/except blackjax import.

73f9e5c

ricardoV94 reviewed Sep 9, 2022

View reviewed changes

pymc_experimental/inference/fit.py Show resolved Hide resolved

pymc_experimental/inference/fit.py Outdated Show resolved Hide resolved

twiecki and others added 5 commits September 9, 2022 16:20

Update pymc_experimental/inference/fit.py

cf3d0de

Co-authored-by: Ricardo Vieira <[email protected]>

Update pymc_experimental/inference/fit.py

6deca7b

Co-authored-by: Ricardo Vieira <[email protected]>

Move blackjax to dev reqs.

0bbfb56

Make import non-optional.

7f7fcb3

Precommit.

e2dad01

ricardoV94 requested changes Sep 9, 2022

View reviewed changes

twiecki added 5 commits September 9, 2022 16:30

Change imports.

4fc5e89

Call fit() from test.

7bdb10a

Fix fit import.

8f07c25

Fix kwargs.

350b77d

Add blackjax to test env.

7734b03

ricardoV94 changed the title ~~Add wrapper for running blackjax pathfinder.~~ Add wrapper for running blackjax pathfinder Sep 9, 2022

twiecki added 3 commits September 9, 2022 17:17

Only look for tests in test subdir.

625caf0

Only look for tests in test subdir.

2d0a435

Only look for tests in test subdir.

0dca76c

ricardoV94 approved these changes Sep 9, 2022

View reviewed changes

ricardoV94 merged commit 75aa37b into main Sep 9, 2022

twiecki deleted the blackjax_pathfinder branch September 9, 2022 15:32

twiecki mentioned this pull request Sep 10, 2022

Broken CI #76

Closed

twiecki restored the blackjax_pathfinder branch September 10, 2022 09:27

twiecki added a commit that referenced this pull request Sep 10, 2022

Revert "Add wrapper for running blackjax pathfinder (#72)"

bf6b4d3

This reverts commit 75aa37b.

twiecki mentioned this pull request Sep 10, 2022

Revert "Add wrapper for running blackjax pathfinder" #78

Merged

twiecki added a commit that referenced this pull request Sep 10, 2022

Revert "Add wrapper for running blackjax pathfinder (#72)" (#78)

9924d92

This reverts commit 75aa37b.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add wrapper for running blackjax pathfinder #72

Add wrapper for running blackjax pathfinder #72

twiecki commented Sep 7, 2022 •

edited

Loading

twiecki commented Sep 8, 2022

twiecki commented Sep 8, 2022

ricardoV94 left a comment

ricardoV94 commented Sep 8, 2022

ricardoV94 commented Sep 8, 2022

twiecki commented Sep 8, 2022

ColCarroll commented Sep 8, 2022

ColCarroll commented Sep 8, 2022

rlouf commented Sep 8, 2022

ricardoV94 commented Sep 8, 2022

ricardoV94 commented Sep 9, 2022

twiecki commented Sep 9, 2022

ricardoV94 left a comment

ricardoV94 Sep 9, 2022

ricardoV94 commented Sep 10, 2022 •

edited

Loading

twiecki commented Sep 10, 2022

Add wrapper for running blackjax pathfinder #72

Add wrapper for running blackjax pathfinder #72

Conversation

twiecki commented Sep 7, 2022 • edited Loading

twiecki commented Sep 8, 2022

twiecki commented Sep 8, 2022

ricardoV94 left a comment

Choose a reason for hiding this comment

ricardoV94 commented Sep 8, 2022

ricardoV94 commented Sep 8, 2022

twiecki commented Sep 8, 2022

ColCarroll commented Sep 8, 2022

ColCarroll commented Sep 8, 2022

rlouf commented Sep 8, 2022

ricardoV94 commented Sep 8, 2022

ricardoV94 commented Sep 9, 2022

twiecki commented Sep 9, 2022

ricardoV94 left a comment

Choose a reason for hiding this comment

ricardoV94 Sep 9, 2022

Choose a reason for hiding this comment

ricardoV94 commented Sep 10, 2022 • edited Loading

twiecki commented Sep 10, 2022

twiecki commented Sep 7, 2022 •

edited

Loading

ricardoV94 commented Sep 10, 2022 •

edited

Loading