Access to external nuts samplers via `sample` #6422

twiecki · 2023-01-01T06:07:54Z

TODO:

Tests
Figure out circular dependency on nutpie test...

Major

All arguments except draws in sample are now keyword-only.

Enhancements

Access to external nuts samplers via sample

codecov · 2023-01-01T06:21:31Z

Codecov Report

Merging #6422 (f489848) into main (301b28a) will decrease coverage by 0.01%.
The diff coverage is 91.80%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6422      +/-   ##
==========================================
- Coverage   94.74%   94.74%   -0.01%     
==========================================
  Files         146      147       +1     
  Lines       27807    27861      +54     
==========================================
+ Hits        26346    26397      +51     
- Misses       1461     1464       +3

Impacted Files	Coverage Δ
pymc/tests/distributions/test_mixture.py	`99.35% <ø> (ø)`
pymc/sampling/mcmc.py	`92.28% <82.14%> (-0.75%)`	⬇️
pymc/sampling/jax.py	`98.26% <100.00%> (+<0.01%)`	⬆️
pymc/tests/sampling/test_jax.py	`100.00% <100.00%> (ø)`
pymc/tests/sampling/test_mcmc_external.py	`100.00% <100.00%> (ø)`
pymc/model.py	`90.13% <0.00%> (+0.27%)`	⬆️

pymc/sampling/mcmc.py

ferrine · 2023-01-01T12:22:11Z

So far nutpie doesn’t provide precompiled Python wheels. So users that pip install pymc are unable to use this sampler. pymc-devs/nutpie#22

ricardoV94 · 2023-01-01T12:43:29Z

So far nutpie doesn’t provide precompiled Python wheels. So users that pip install pymc are unable to use this sampler. pymc-devs/nutpie#22

That's okay, we're not making it the default here. There's a try except import with an informative message

fonnesbeck · 2023-01-02T03:00:51Z

nutpie can currently only be installed via maturin, correct?

twiecki · 2023-01-02T05:17:52Z

nutpie can currently only be installed via maturin, correct?

It's also available via conda-forge.

pymc/sampling/mcmc.py

twiecki · 2023-01-03T09:04:33Z

[pymc/sampling/mcmc.py]
pymc/sampling/mcmc.py:253: error: Argument "chains" to "sample_numpyro_nuts" has incompatible type "Optional[int]"; expected "int"
pymc/sampling/mcmc.py:255: error: Argument "random_seed" to "sample_numpyro_nuts" has incompatible type "Union[Union[int, Sequence[int], ndarray[Any, Any], None], RandomState, Generator, None]"; expected "Union[int, Sequence[int], ndarray[Any, Any], None]"
pymc/sampling/mcmc.py:266: error: Argument "chains" to "sample_blackjax_nuts" has incompatible type "Optional[int]"; expected "int"
pymc/sampling/mcmc.py:268: error: Argument "random_seed" to "sample_blackjax_nuts" has incompatible type "Union[Union[int, Sequence[int], ndarray[Any, Any], None], RandomState, Generator, None]"; expected "Union[int, Sequence[int], ndarray[Any, Any], None]"

Is this important? Not very familiar with types.

ricardoV94 · 2023-02-02T18:31:27Z

Hmm a specific version of pymc is being installed by nutpie... so the test doesn't make much sense...

pymc/sampling/mcmc.py

twiecki · 2023-02-13T13:35:04Z

@ricardoV94 Any progress? What's the status? Should someone else take over?

ricardoV94 · 2023-02-13T14:17:11Z

The problem is I don't know how to test nutpie, because nutpie has a direct dependency in pymc/pytensor (at least the conda package). I'll see if it can be pip installed instead. But it may not make sense to test at all given the circular dependency

twiecki · 2023-02-13T15:48:44Z

OK, lets just not test it then, or only if it's istalled.

twiecki · 2023-02-13T18:18:25Z

Tests pass! Last comment but LGTM otherwise.

ricardoV94 · 2023-02-13T18:29:26Z

Tests already passed before, they were failing for unrelated reasons :)

Still didn't remove the test dependency on nutpie

pymc/sampling/mcmc.py

michaelosthege · 2023-02-13T20:22:47Z

pymc/sampling/mcmc.py

@@ -503,7 +613,7 @@ def sample(
    )

    sample_args = {
-        "draws": draws,
+        "draws": draws + tune,  # FIXME: Why is tune added to draws?


probably a relic of old signatures

I will open an issue after this PR to investigate it

michaelraczycki · 2023-02-14T09:17:24Z

Have you tried using dependency injection? Because if yes, other way would be to use unittest.mock, I've used it in the past but that will depend on how complex the test is.

ricardoV94 · 2023-02-14T10:17:53Z

Maybe the dependency is not too bad. It first installs nutpie and then the local pymc, so I guess if PyMC depends on a newer PyTensor it will replace it by that? We can merge and see what happens when we next bump the PyTensor dependency.

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

Reorder arguments more logically [citation needed]

twiecki · 2023-02-15T19:17:38Z

This looks great! merge?

ricardoV94 · 2023-02-15T20:22:22Z

This looks great! merge?

Never hurts to see what happens next :D

twiecki · 2023-02-15T20:25:55Z

Let's do a minor release!

ricardoV94 reviewed Jan 1, 2023

View reviewed changes

pymc/sampling/mcmc.py Show resolved Hide resolved

ricardoV94 reviewed Jan 2, 2023

View reviewed changes

pymc/sampling/mcmc.py Outdated Show resolved Hide resolved

pymc/sampling/mcmc.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jan 3, 2023

View reviewed changes

pymc/sampling/mcmc.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jan 3, 2023

View reviewed changes

pymc/sampling/mcmc.py Show resolved Hide resolved

twiecki changed the title ~~Add fast_nuts option to sample() that tries different compiled samplers.~~ Add compiled nuts samplers to sample() Jan 24, 2023

twiecki assigned ricardoV94 Jan 24, 2023

ricardoV94 force-pushed the fast_nuts branch from 7036682 to 6c8c5ce Compare January 26, 2023 12:32

ricardoV94 changed the title ~~Add compiled nuts samplers to sample()~~ Add access to external nuts samplers via sample Jan 26, 2023

ricardoV94 added the enhancements label Jan 26, 2023

ricardoV94 force-pushed the fast_nuts branch from 6c8c5ce to 94d1ebe Compare January 26, 2023 12:34

ricardoV94 marked this pull request as draft January 26, 2023 12:36

ricardoV94 changed the title ~~Add access to external nuts samplers via sample~~ Access to external nuts samplers via sample Jan 26, 2023

ricardoV94 force-pushed the fast_nuts branch 3 times, most recently from f2dfcaa to 5b5b9c9 Compare February 2, 2023 17:56

ricardoV94 marked this pull request as ready for review February 2, 2023 18:28

ricardoV94 requested review from cluhmann, michaelosthege and aseyboldt and removed request for cluhmann February 2, 2023 18:31

twiecki commented Feb 3, 2023

View reviewed changes

pymc/sampling/mcmc.py Outdated Show resolved Hide resolved

ricardoV94 mentioned this pull request Feb 3, 2023

Allow JAX samplers when available. pymc-labs/pymc-marketing#144

Closed

ricardoV94 force-pushed the fast_nuts branch 2 times, most recently from f88fadb to 27bc414 Compare February 13, 2023 16:59

ricardoV94 force-pushed the fast_nuts branch from 27bc414 to c18522f Compare February 13, 2023 18:41

michaelosthege reviewed Feb 13, 2023

View reviewed changes

ricardoV94 and others added 3 commits February 15, 2023 15:42

Remove experimental module warning from jax.py

f108dcf

Raise error when trying to jaxify graphs with RNGs

e433b76

Allow external nuts sampler directly from sample

4e4d234

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

ricardoV94 force-pushed the fast_nuts branch from c18522f to 8df3fec Compare February 15, 2023 14:42

ricardoV94 added the major Include in major changes release notes section label Feb 15, 2023

Make all but draws keyword-only arguments in sample.

f489848

Reorder arguments more logically [citation needed]

ricardoV94 force-pushed the fast_nuts branch from 8df3fec to f489848 Compare February 15, 2023 15:36

twiecki merged commit c2ce47f into main Feb 15, 2023

twiecki deleted the fast_nuts branch February 15, 2023 20:25

twiecki mentioned this pull request Mar 6, 2023

New nuts_sampler kwarg for sample() pymc-devs/communication#2

Closed

Uh oh!

Access to external nuts samplers via sample #6422

Access to external nuts samplers via sample #6422

Uh oh!

Conversation

twiecki commented Jan 1, 2023 • edited by ricardoV94 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Major

Enhancements

Uh oh!

codecov bot commented Jan 1, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

ferrine commented Jan 1, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 commented Jan 1, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

fonnesbeck commented Jan 2, 2023

Uh oh!

twiecki commented Jan 2, 2023

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

twiecki commented Jan 3, 2023

Uh oh!

ricardoV94 commented Feb 2, 2023

Uh oh!

Uh oh!

twiecki commented Feb 13, 2023

Uh oh!

ricardoV94 commented Feb 13, 2023

Uh oh!

twiecki commented Feb 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

twiecki commented Feb 13, 2023

Uh oh!

ricardoV94 commented Feb 13, 2023

Uh oh!

Uh oh!

michaelosthege Feb 13, 2023

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Feb 14, 2023

Choose a reason for hiding this comment

Uh oh!

michaelraczycki commented Feb 14, 2023

Uh oh!

ricardoV94 commented Feb 14, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

twiecki commented Feb 15, 2023

Uh oh!

ricardoV94 commented Feb 15, 2023

Uh oh!

twiecki commented Feb 15, 2023

Uh oh!

Uh oh!

Access to external nuts samplers via `sample` #6422

Access to external nuts samplers via `sample` #6422

twiecki commented Jan 1, 2023 •

edited by ricardoV94

Loading

codecov bot commented Jan 1, 2023 •

edited

Loading

ferrine commented Jan 1, 2023 •

edited

Loading

ricardoV94 commented Jan 1, 2023 •

edited

Loading

twiecki commented Feb 13, 2023 •

edited

Loading

ricardoV94 commented Feb 14, 2023 •

edited

Loading