fix: restore blackjax>=1.6 compatibility for nuts_sampler="blackjax" - #8373
fix: restore blackjax>=1.6 compatibility for nuts_sampler="blackjax"#8373laishettikarthik-tech wants to merge 13 commits into
Conversation
Handle progress_bar parameter for blackjax version compatibility.
Add regression tests for blackjax progress bar compatibility.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #8373 +/- ##
==========================================
- Coverage 91.83% 90.23% -1.60%
==========================================
Files 128 128
Lines 21259 21268 +9
==========================================
- Hits 19523 19192 -331
- Misses 1736 2076 +340
🚀 New features to boost your workflow:
|
|
pre-commit.ci autofix |
for more information, see https://pre-commit.ci
notluquis
left a comment
notluquis
left a comment
There was a problem hiding this comment.
Confirming the reproduction on a configuration not listed in #8367: blackjax 1.6.2, pymc 6.1.0, jax 0.10.2, Python 3.13, macOS arm64. The ordering bug is also present in the released 6.2.0 source, not only on main, so upgrading is not a workaround.
Applied this branch locally with the two changes suggested below. The issue's own reproduction then runs on both paths:
progressbar=False -> OK, posterior shape (1, 50), mean 0.1311
progressbar=True -> OK, warns about the missing extra and samples without the bar
The hasattr(blackjax.progress_bar, "gen_scan_fn") dispatch checks out: on 1.6.2 blackjax.progress_bar is a function rather than a module, hasattr(..., "gen_scan_fn") is False, and its signature is (label='BlackJAX', print_rate=None, output_file=None), so label="NUTS" is valid.
Separate thought, out of scope for this fix. The reason this surfaced as TypeError: build_kernel.<locals>.kernel() got an unexpected keyword argument — naming a function nobody called, from inside JAX tracing — is that window_adaptation forwards unknown **extra_parameters straight through to the algorithm. Filtering against the real signature would turn any future stale kwarg into an error at its origin:
accepted = inspect.signature(blackjax.window_adaptation).parameters
unknown = set(adaptation_kwargs) - set(accepted)
if unknown:
raise TypeError(f"blackjax.window_adaptation does not accept {sorted(unknown)}")A test for progress_bar specifically would not catch the next blackjax API change; this would.
| elif progress_bar: | ||
| # blackjax >= 1.6: progress_bar is a context manager that | ||
| # monkeypatches jax.lax.scan for its duration instead. | ||
| with blackjax.progress_bar(label="NUTS"): | ||
| _, (samples, stats) = jax.lax.scan(_one_step, last_state, (jnp.arange(draws), keys)) |
There was a problem hiding this comment.
On blackjax >= 1.6 this context manager is powered by jaxtap, shipped as blackjax[progress]. On a plain blackjax==1.6.2 install it raises:
ImportError: blackjax.progress_bar requires the 'progress' optional extra.
Install it with: pip install 'blackjax[progress]'
So test_sample_blackjax_nuts_progressbar_true will fail on any CI image that installs blackjax without the extra, and a user passing progressbar=True trades a TypeError for an ImportError. A progress bar aborting a sampling run seems worse than losing the bar, so a fallback may be preferable to adding the extra as a test dependency (warnings is already imported at line 17):
| elif progress_bar: | |
| # blackjax >= 1.6: progress_bar is a context manager that | |
| # monkeypatches jax.lax.scan for its duration instead. | |
| with blackjax.progress_bar(label="NUTS"): | |
| _, (samples, stats) = jax.lax.scan(_one_step, last_state, (jnp.arange(draws), keys)) | |
| elif progress_bar: | |
| try: | |
| with blackjax.progress_bar(label="NUTS"): | |
| _, (samples, stats) = jax.lax.scan(_one_step, last_state, (jnp.arange(draws), keys)) | |
| except ImportError: | |
| warnings.warn( | |
| "blackjax progress bar needs the 'progress' extra " | |
| "(pip install 'blackjax[progress]'); sampling without it.", | |
| UserWarning, | |
| stacklevel=2, | |
| ) | |
| _, (samples, stats) = jax.lax.scan(_one_step, last_state, (jnp.arange(draws), keys)) |
| return samples, stats | ||
|
|
||
| keys = jax.random.split(seed, draws) | ||
| if hasattr(blackjax.progress_bar, "gen_scan_fn"): | ||
| # blackjax < 1.6: progress_bar is a module exposing gen_scan_fn, | ||
| # which wraps jax.lax.scan directly. | ||
| scan_fn = blackjax.progress_bar.gen_scan_fn(draws, progress_bar) | ||
| _, (samples, stats) = scan_fn(_one_step, last_state, (jnp.arange(draws), keys)) | ||
| elif progress_bar: | ||
| # blackjax >= 1.6: progress_bar is a context manager that | ||
| # monkeypatches jax.lax.scan for its duration instead. | ||
| with blackjax.progress_bar(label="NUTS"): | ||
| _, (samples, stats) = jax.lax.scan(_one_step, last_state, (jnp.arange(draws), keys)) | ||
| else: | ||
| _, (samples, stats) = jax.lax.scan(_one_step, last_state, (jnp.arange(draws), keys)) |
There was a problem hiding this comment.
This block repeats lines 283-295 verbatim and sits after the return samples, stats above it, so it is unreachable. Looks like a rebase artefact. It also accounts for the Codecov number — 55.6% patch coverage is roughly what you get when half the added lines cannot execute.
Deleting the added copy leaves the original return samples, stats below it, so the diff does not touch any pre-existing line.
| return samples, stats | |
| keys = jax.random.split(seed, draws) | |
| if hasattr(blackjax.progress_bar, "gen_scan_fn"): | |
| # blackjax < 1.6: progress_bar is a module exposing gen_scan_fn, | |
| # which wraps jax.lax.scan directly. | |
| scan_fn = blackjax.progress_bar.gen_scan_fn(draws, progress_bar) | |
| _, (samples, stats) = scan_fn(_one_step, last_state, (jnp.arange(draws), keys)) | |
| elif progress_bar: | |
| # blackjax >= 1.6: progress_bar is a context manager that | |
| # monkeypatches jax.lax.scan for its duration instead. | |
| with blackjax.progress_bar(label="NUTS"): | |
| _, (samples, stats) = jax.lax.scan(_one_step, last_state, (jnp.arange(draws), keys)) | |
| else: | |
| _, (samples, stats) = jax.lax.scan(_one_step, last_state, (jnp.arange(draws), keys)) |
Refactor progress bar handling for blackjax sampling.
|
The tests / all_tests failure is unrelated to this PR — it's test_mvstudentt[NUMBA] in test_random_alternative_backends.py, a numerical tolerance mismatch in the Numba multivariate Student-T sampler, nothing to do with jax.py or blackjax. All 133 other tests passed, including the three TestBlackjaxProgressBarCompat tests. Ran the blackjax-specific tests locally too (see earlier comment) — all pass cleanly. Given it's unrelated, might this need a rerun, or is test_mvstudentt[NUMBA] a known flake? |
| import blackjax | ||
|
|
||
| from pymc.sampling.jax import _blackjax_inference_loop |
There was a problem hiding this comment.
check if local imports can be made global
Refactor test for progress bar handling in JAX sampling.
…ing import, add trailing newline
3 participants
blackjax 1.6 removed the progress_bar parameter from window_adaptation
(unknown kwargs now forward into the NUTS kernel and raise TypeError)
and replaced the progress_bar module (gen_scan_fn helper) with a
context-manager function.
window_adaptation, not after
both the old module-based API and the new context-manager API
Closes #8367