Adding NUTS sampler from blackjax to sampling_jax #5477

zaxtax · 2022-02-16T02:53:39Z

This PR adds blackjax support into sampling_jax

It's not perfect (progress bar support and sampling stats are not fully supported in blackjax). But it works with all transformed and deterministic variables.

are the changes—especially new features—covered by tests and docstrings?
linting/style checks have been run

zaxtax · 2022-02-16T02:55:20Z

Although not mentioned @junpenglao and @rlouf helped a lot with getting this to work from the blackjax side.

codecov · 2022-02-16T03:03:36Z

Codecov Report

Merging #5477 (d5fcbb0) into main (44c5495) will increase coverage by 0.05%.
The diff coverage is 94.73%.

@@            Coverage Diff             @@
##             main    #5477      +/-   ##
==========================================
+ Coverage   87.59%   87.64%   +0.05%     
==========================================
  Files          76       76              
  Lines       13694    13765      +71     
==========================================
+ Hits        11995    12065      +70     
- Misses       1699     1700       +1

Impacted Files	Coverage Δ
pymc/sampling_jax.py	`96.96% <94.73%> (-1.46%)`	⬇️
pymc/parallel_sampling.py	`87.70% <0.00%> (+0.99%)`	⬆️

rlouf · 2022-02-16T04:38:28Z

Great! We can add a progress with callbacks, as is done in MCX. Afair the overhead it introduces is very small.

ricardoV94

Looks great. I'll have another look later, but left a comment above

pymc/sampling_jax.py

twiecki · 2022-02-16T10:23:19Z

This is awesome, thanks @zaxtax!
Closes #5454

pymc/sampling_jax.py

ricardoV94 · 2022-02-16T15:04:16Z

pymc/sampling_jax.py

+    model_logpt = model.logpt()
+    if not negative_logp:
+        model_logpt = -model_logpt
+    logp_fn = get_jaxified_graph(inputs=model.value_vars, outputs=[model_logpt])

    def logp_fn_wrap(x):
        # NumPyro expects a scalar potential with the opposite sign of model.logpt
        res = logp_fn(*x)[0]


Let's return this directly, instead of saving it intermediately in res

OriolAbril · 2022-02-16T16:01:00Z

pymc/sampling_jax.py

+    return inference_loop(seed, last_state)
+
+
+def sample_blackjax_nuts(


This function needs to be documented and added to the API docs (where the numpyro one might also be missing?). It probably makes more sense to add in the first block of https://github.com/pymc-devs/pymc/blob/main/docs/source/api/samplers.rst.

I'm a bit rusty with sphinx, where do I specify that the functions are in sampling_jax

by that point currentmodule is set to pymc, so sphinx assumes functions live there. If they are in pymc.sampling_jax.<function> you'll need to use sampling_jax.<function> as the pymc is assumed because of the currentmodule use

OriolAbril · 2022-02-16T16:01:55Z

pymc/sampling_jax.py

+        observed_data=find_observations(model),
+        coords=coords,
+        dims=dims,
+        attrs={"sampling_time": (tic3 - tic2).total_seconds()},


can we also add pymc and blackjax versions?

Sure thing!

Oh is there an arviz convention for when there are multiple libraries? This came up in earlier discussions.

not that I'm aware of 🤔

I'm going to put in blackjax and numpyro as the inference library, and add pymc_version as another field.

OriolAbril · 2022-02-16T18:38:35Z

docs/source/api/samplers.rst

@@ -30,6 +30,8 @@ HMC family

   NUTS
   HamiltonianMC
+   sampling_jax.sample_blackjax_nuts


I think they should be in the first block (or in a different subsection), IIUC, both are nuts only even if blackjax has other samplers, but they are not step methods that can be used as inputs to pm.sample

OriolAbril · 2022-02-16T18:45:30Z

pymc/sampling_jax.py

@@ -188,6 +188,45 @@ def sample_blackjax_nuts(
    chain_method="parallel",
    idata_kwargs=None,
 ):
+    """Draw samples from the posterior using the NUTS method from the blackjax library.
+
+    Parameters


Can you update with the advise on https://pymc-data-umbrella.xyz/en/latest/sprint/docstring_tutorial.html#edit-the-docstring and I'll take a 2nd look later?

I must be missing something but I think I'm style compliant. I actually copied the docstrings from sample as my starting point.

Basically none of the docstrings are currently style compliant, and most aren't rendered correctly either because of that. We have an issue open for this: #5459. I can give a quick go at one of the docstrings. I think you'll also need to rebase on main for the docs preview to work and CI to pass.

added some comments, but they are not exaustive

Oh ok it was how I'm specifying defaults and optional.

OriolAbril · 2022-02-16T19:53:11Z

pymc/sampling_jax.py

+    draws : int
+        The number of samples to draw. Defaults to 1000. The number of tuned samples are discarded
+        by default.


Suggested change

draws : int

The number of samples to draw. Defaults to 1000. The number of tuned samples are discarded

by default.

draws : int, default 1000

The number of samples to draw. The number of tuned samples are discarded

by default.

OriolAbril · 2022-02-16T19:55:46Z

pymc/sampling_jax.py

+    model : Model (optional if in ``with`` context)
+        Model to sample from. The model needs to have free random variables.


Suggested change

model : Model (optional if in ``with`` context)

Model to sample from. The model needs to have free random variables.

model : Model, optional

Model to sample from. The model needs to have free random variables. When inside a ``with`` model

context, it defaults to that model, otherwise the model must be passed explicitly.

OriolAbril · 2022-02-16T19:56:20Z

pymc/sampling_jax.py

+    var_names : Iterable[str]
+        Names of variables for which to compute the posterior samples.


Suggested change

var_names : Iterable[str]

Names of variables for which to compute the posterior samples.

var_names : iterable of str, optional

Names of variables for which to compute the posterior samples. Defaults to all variables in the posterior

OriolAbril · 2022-02-16T19:56:42Z

pymc/sampling_jax.py

+        Model to sample from. The model needs to have free random variables.
+    var_names : Iterable[str]
+        Names of variables for which to compute the posterior samples.
+    progress_bar : bool, optional default=True


Suggested change

progress_bar : bool, optional default=True

progress_bar : bool, default True

OriolAbril · 2022-02-16T19:57:04Z

pymc/sampling_jax.py

+        Specify how samples should be drawn. The choices include "parallel", and
+        "vectorized". Defaults to "parallel".
+    idata_kwargs : dict, optional
+        Keyword arguments for :func:`pymc.to_inference_data`


Suggested change

Keyword arguments for :func:`pymc.to_inference_data`

Keyword arguments for :func:`arviz.from_dict`. It also accepts a boolean as value for the ``log_likelihood``

key to indicate that the pointwise log likelihood should not be included in the returned object.

OriolAbril · 2022-02-16T19:58:20Z

pymc/sampling_jax.py

+    trace : arviz.InferenceData
+        ArviZ ``InferenceData`` object that contains the samples.


Suggested change

trace : arviz.InferenceData

ArviZ ``InferenceData`` object that contains the samples.

InferenceData

ArviZ ``InferenceData`` object that contains the posterior samples, together with their respective sample stats and

pointwise log likeihood values (unless skipped with ``idata_kwargs``).

OriolAbril · 2022-02-16T19:59:48Z

pymc/sampling_jax.py

+    tune : int
+        Number of iterations to tune, defaults to 1000. Samplers adjust the step sizes, scalings or


Suggested change

tune : int

Number of iterations to tune, defaults to 1000. Samplers adjust the step sizes, scalings or

tune : int, default 1000

Number of iterations to tune. Samplers adjust the step sizes, scalings or

OriolAbril · 2022-02-16T21:47:40Z

The functions are still missing from docs (https://pymc--5477.org.readthedocs.build/en/5477/api/samplers.html) due to this issue:

WARNING: [autosummary] failed to import pymc.sampling_jax.sample_blackjax_nuts.
Possible hints:
* AttributeError: module 'pymc' has no attribute 'sampling_jax'
* ImportError: no module named pymc.sampling_jax
* ModuleNotFoundError: No module named 'jax'
WARNING: [autosummary] failed to import pymc.sampling_jax.sample_numpyro_nuts.
Possible hints:
* AttributeError: module 'pymc' has no attribute 'sampling_jax'
* ImportError: no module named pymc.sampling_jax
* ModuleNotFoundError: No module named 'jax'

what do we think is a better option, should jax be imported in a try except or installed in the docs env?

zaxtax · 2022-02-16T21:53:29Z

The functions are still missing from docs (https://pymc--5477.org.readthedocs.build/en/5477/api/samplers.html) due to this issue:
WARNING: [autosummary] failed to import pymc.sampling_jax.sample_blackjax_nuts.
Possible hints:
* AttributeError: module 'pymc' has no attribute 'sampling_jax'
* ImportError: no module named pymc.sampling_jax
* ModuleNotFoundError: No module named 'jax'
WARNING: [autosummary] failed to import pymc.sampling_jax.sample_numpyro_nuts.
Possible hints:
* AttributeError: module 'pymc' has no attribute 'sampling_jax'
* ImportError: no module named pymc.sampling_jax
* ModuleNotFoundError: No module named 'jax'
what do we think is a better option, should jax be imported in a try except or installed in the docs env?

I'm already installing jax in the other test environment. But I don't maybe there should be a try-except in the module itself. Would the latter suppress this error?

OriolAbril · 2022-02-16T22:15:22Z

yep, both should fix the issue. Sphinx imports the objects to be documented in order to extract the __doc__ attributes, but it doesn't use them. The problem here is that it errors out before being able to import the two functions due to the jax import. I don't know the best way to go about it because I don't really understand how dependencies+dev+test+optional dependencies are organized or how problematic it would be to modify this import (I haven't used jax samplers yet)

zaxtax · 2022-02-16T22:27:23Z

yep, both should fix the issue. Sphinx imports the objects to be documented in order to extract the __doc__ attributes, but it doesn't use them. The problem here is that it errors out before being able to import the two functions due to the jax import. I don't know the best way to go about it because I don't really understand how dependencies+dev+test+optional dependencies are organized or how problematic it would be to modify this import (I haven't used jax samplers yet)

As the name implies, jax is required to really use anything in this module. But I think we can suppress some of this stuff.

twiecki · 2022-02-19T19:41:24Z

Aren't the two almost identical except for compilers and fixing blas?

…

On Sat, Feb 19, 2022, 13:31 Ricardo Vieira ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In conda-envs/environment-dev-py37.yml <#5477 (comment)>: > @@ -13,6 +13,7 @@ dependencies: - fastprogress>=0.2.0 - h5py>=2.7 - ipython>=7.16 +- jax I think docs requires quite more dependencies than tests — Reply to this email directly, view it on GitHub <#5477 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAFETGD4JPUWRWQSCIQPCWTU36EQPANCNFSM5OQKUNPQ> . Triage notifications on the go with GitHub Mobile for iOS <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675> or Android <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub>. You are receiving this because you commented.Message ID: ***@***.***>

OriolAbril · 2022-02-20T02:46:13Z

Aren't the two almost identical except for compilers and fixing blas?

assuming this is in reference to test and docs environments. Both have the common step of installing pymc but that is the only thing they have in common. Test then needs to execute and test virtually all the functions so it needs all optional dependencies (even runtime ones), needs pytest, things for lint, pre-commit, coverage... Whereas docs only needs to execute 4 notebooks, for everything else importing the objects is enough, and it needs sphinx and all the extensions which explode in dependencies very quickly, myst-nb needs myst-parser and jupyter to execute and render notebooks for example.

twiecki · 2022-02-28T12:33:50Z

pymc/sampling_jax.py

@@ -457,7 +462,7 @@ def sample_numpyro_nuts(

    if random_seed is None:
        random_seed = model.rng_seeder.randint(
-            2**30, dtype=np.int64, size=chains if chains > 1 else None
+            2 ** 30, dtype=np.int64, size=chains if chains > 1 else None


I think you need to update pre-commit.

How do I do that?

zaxtax · 2022-02-28T12:54:35Z

I don't know what's going on with the tests failing now, but this runs locally for me.

I've rebased from main, but would like this go in sooner since it's a larger contribution.

Postponing adding progress bar support until callback support has been included within blackjax

ericmjl · 2022-02-28T13:43:34Z

@zaxtax curious to hear, in reading the logs, I see a "core dumped" error showing up. Could the issue be RAM usage?

The GH Actions runners have ~7GB of RAM, I think, I read here.

If I'm wrong, please feel free to correct me!

zaxtax · 2022-03-03T15:04:02Z

@ricardoV94 any idea why testing is timing out?

ricardoV94 · 2022-03-03T15:20:13Z

@ricardoV94 any idea why testing is timing out?

Nope

zaxtax · 2022-03-03T16:17:05Z

@ricardoV94 any idea why testing is timing out?

Nope

Then I'm a bit out of ideas how to push this PR forward.

zaxtax · 2022-03-10T15:58:02Z

@ricardoV94 looks like the tests pass now!

ricardoV94 · 2022-03-10T19:34:03Z

@ricardoV94 looks like the tests pass now!

Great! I manually made the docs rebuild to see if the failure was just flaky

zaxtax · 2022-03-10T21:22:36Z

@ricardoV94 looks like the tests pass now!

Great! I manually made the docs rebuild to see if the failure was just flaky

Looks like it does!

ericmjl · 2022-03-10T23:22:55Z

Looks like we’re good to go!

zaxtax · 2022-03-11T20:14:26Z

@ricardoV94 is this ready to merge?

twiecki · 2022-03-11T20:18:49Z

Thanks @zaxtax, awesome to have this in!

twiecki · 2022-03-11T20:20:40Z

One thing we forgot are release notes.

zaxtax · 2022-03-11T21:39:43Z

I haven't done that for a lot of my `sampling_jax` contributions

…

On Fri, 11 Mar 2022, 21:20 Thomas Wiecki, ***@***.***> wrote: One thing we forgot are release notes. — Reply to this email directly, view it on GitHub <#5477 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCULYOC4ZCZ7XE62YSSTU7OTKJANCNFSM5OQKUNPQ> . Triage notifications on the go with GitHub Mobile for iOS <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675> or Android <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub>. You are receiving this because you were mentioned.Message ID: ***@***.***>

ricardoV94 · 2022-03-12T08:33:36Z

I haven't done that for a lot of my sampling_jax contributions
…
On Fri, 11 Mar 2022, 21:20 Thomas Wiecki, @.> wrote: One thing we forgot are release notes. — Reply to this email directly, view it on GitHub <#5477 (comment)>, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAACCULYOC4ZCZ7XE62YSSTU7OTKJANCNFSM5OQKUNPQ . Triage notifications on the go with GitHub Mobile for iOS https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675 or Android https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub. You are receiving this because you were mentioned.Message ID: @.>

Yeah we were in this regime where sampling_jax is still "experimental" but now with it showing on the docs and all the work we have been doing is getting more and more official. So we should try to remember and mention changes in release notes, specially these large ones!

We definitely need one for this

ricardoV94 reviewed Feb 16, 2022

View reviewed changes

pymc/sampling_jax.py Outdated Show resolved Hide resolved

twiecki reviewed Feb 16, 2022

View reviewed changes

pymc/sampling_jax.py Show resolved Hide resolved

ricardoV94 added the jax label Feb 16, 2022

ricardoV94 reviewed Feb 16, 2022

View reviewed changes

pymc/sampling_jax.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Feb 16, 2022

View reviewed changes

OriolAbril reviewed Feb 16, 2022

View reviewed changes

zaxtax force-pushed the adding-blackjax-support branch from e1a636f to 8781b32 Compare February 16, 2022 20:35

zaxtax added 9 commits February 16, 2022 20:38

Adding NUTS sampler from blackjax to sampling_jax

5a36685

Lint fixes

41425fb

Refactor get_jaxified_logp

1a9e8d6

Install blackjax in workflows

ec62c89

Fix url

2394537

Simplify function

0e01e26

Add documentation

75df8d2

Add library versions

3b5c948

Move to more appropriate section

8a27c1f

zaxtax force-pushed the adding-blackjax-support branch from 8781b32 to 8a27c1f Compare February 16, 2022 20:38

Fix docstrings

ea5874b

zaxtax added 2 commits February 16, 2022 22:59

Add jax to dev environment

659828c

Jax doesn't work on windows, remove it

6460a16

twiecki approved these changes Feb 20, 2022

View reviewed changes

zaxtax added 2 commits February 28, 2022 12:17

Merge branch 'main' into adding-blackjax-support

df86b55

Fix merge

6ba2c31

twiecki reviewed Feb 28, 2022

View reviewed changes

zaxtax added 2 commits February 28, 2022 12:34

Remove progress bar functionality for now

1bf85f0

Pre-commit fix

b1436d9

twiecki requested a review from ricardoV94 February 28, 2022 16:20

Merge branch 'main' into adding-blackjax-support

c2a5ea7

ricardoV94 added this to the v4.0.0b4 milestone Mar 7, 2022

Merge branch 'main' into adding-blackjax-support

d5fcbb0

twiecki merged commit b799547 into pymc-devs:main Mar 11, 2022

ricardoV94 mentioned this pull request Mar 12, 2022

Mention blackjax nuts sampler in release notes #5587

Closed

ricardoV94 mentioned this pull request Mar 19, 2022

Add blackjax sampler to sampling_jax.py #5454

Closed

isms mentioned this pull request Jul 1, 2022

Add numpyro install to building docs instructions #5936

Merged

		return inference_loop(seed, last_state)


		def sample_blackjax_nuts(

		model : Model (optional if in ``with`` context)
		Model to sample from. The model needs to have free random variables.

-    model : Model (optional if in ``with`` context)
-        Model to sample from. The model needs to have free random variables.
+    model : Model, optional
+        Model to sample from. The model needs to have free random variables. When inside a ``with`` model
+        context, it defaults to that model, otherwise the model must be passed explicitly.

		var_names : Iterable[str]
		Names of variables for which to compute the posterior samples.

	progress_bar : bool, optional default=True
	progress_bar : bool, default True

	Keyword arguments for :func:`pymc.to_inference_data`
	Keyword arguments for :func:`arviz.from_dict`. It also accepts a boolean as value for the ``log_likelihood``
	key to indicate that the pointwise log likelihood should not be included in the returned object.

		trace : arviz.InferenceData
		ArviZ ``InferenceData`` object that contains the samples.

		tune : int
		Number of iterations to tune, defaults to 1000. Samplers adjust the step sizes, scalings or

Adding NUTS sampler from blackjax to sampling_jax #5477

Adding NUTS sampler from blackjax to sampling_jax #5477

Conversation

zaxtax commented Feb 16, 2022

zaxtax commented Feb 16, 2022

codecov bot commented Feb 16, 2022 • edited Loading

Codecov Report

rlouf commented Feb 16, 2022

ricardoV94 left a comment

Choose a reason for hiding this comment

twiecki commented Feb 16, 2022

ricardoV94 Feb 16, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zaxtax Feb 16, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

OriolAbril commented Feb 16, 2022

zaxtax commented Feb 16, 2022

OriolAbril commented Feb 16, 2022 • edited Loading

zaxtax commented Feb 16, 2022

twiecki commented Feb 19, 2022 via email

OriolAbril commented Feb 20, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zaxtax commented Feb 28, 2022

ericmjl commented Feb 28, 2022 • edited Loading

zaxtax commented Mar 3, 2022

ricardoV94 commented Mar 3, 2022

zaxtax commented Mar 3, 2022

zaxtax commented Mar 10, 2022

ricardoV94 commented Mar 10, 2022

zaxtax commented Mar 10, 2022

ericmjl commented Mar 10, 2022

zaxtax commented Mar 11, 2022

twiecki commented Mar 11, 2022

twiecki commented Mar 11, 2022

zaxtax commented Mar 11, 2022 via email

ricardoV94 commented Mar 12, 2022 • edited Loading

codecov bot commented Feb 16, 2022 •

edited

Loading

ricardoV94 Feb 16, 2022 •

edited

Loading

zaxtax Feb 16, 2022 •

edited

Loading

OriolAbril commented Feb 16, 2022 •

edited

Loading

ericmjl commented Feb 28, 2022 •

edited

Loading

ricardoV94 commented Mar 12, 2022 •

edited

Loading