Repository navigation
run: blackjax with too few JAX devices runs its chains vectorized (#506) - #507
Merged
Merged
Conversation
pymc maps blackjax's chain_method: parallel onto jax.pmap, one chain per device, with no device check. Importing pymc.sampling.jax forces 100 host devices, but only if the JAX backend has not started; once anything in the process has touched JAX there is one device and two chains crash with "cannot select an axis to squeeze out". The nightly hit this through an earlier test file on the same xdist worker (the locked 3.13 shard failed identically, so no upstream release was involved). run.py now warns and switches blackjax to vectorized when the devices are fewer than the chains. numpyro is left alone: its MCMC already warns and runs the chains sequentially. The flat-raw test samples blackjax vectorized so it no longer depends on which file ran first on its worker. Fixes #506 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #506.
Cause
The nightly failure was not an upstream release. The locked Python 3.13 shard failed with the same error. pymc runs blackjax's
chain_method: parallelthroughjax.pmap, one chain per device, without checking the device count. Importingpymc.sampling.jaxforces 100 host devices throughXLA_FLAGS, but that only works if the JAX backend has not started. When an earlier test file on the same xdist worker has already touched JAX, the process has one device. Two chains then crash withcannot select an axis to squeeze out ... shape=(2, 2). The failure depends on which file runs first on a worker, which is why PR #505's own CI passed.The failure reproduces on the locked JAX 0.11.2 in a plain script. Calling
jax.devices()before importing pymc gives one device, and parallel crashes while vectorized samples. Without that call there are 100 devices and parallel works.Change
run.py: when blackjax would get fewer devices than chains, warn and run the chainsvectorized. numpyro is left alone because its MCMC already warns and runs the chains sequentially.tests/test_flat_logit_raw.py: blackjax samples vectorized, so the test no longer depends on which file ran first on its worker.tests/test_blackjax_chain_method.pycovers the helper.samplers.mddocuments the switch.The new tests and the four-backend flat-raw test pass locally in a session that starts JAX before they run.
🤖 Generated with Claude Code