Add from_blackjax converter#174
Conversation
| @requires("_samples") | ||
| def posterior_to_xarray(self): | ||
| """Convert posterior samples to an xarray Dataset.""" | ||
| assert self._samples is not None |
There was a problem hiding this comment.
| assert self._samples is not None |
The requires decorator will have the method return None without ever callling the method body if the given attribute is None. We already have tests for the decorator and all other converters use them without problem. There is no point in adding these asserts
|
removed the assert statements (the @requires decorator handles that), added |
|
|
||
| from .helpers import importorskip | ||
|
|
||
| blackjax = importorskip("blackjax") |
There was a problem hiding this comment.
The tests are not using blackjax at all. Instead of manually creating fixtures with the outputs that are expected from blackjax, those should be generated with blackjax otherwise if something changes in blackjax the tests will continue to pass even though the converter might end up broken
| def test_from_blackjax_single_chain(single_chain_position, draws): | ||
| # single-chain position should get a chain dim of 1 inserted automatically | ||
| idata = from_blackjax(posterior=single_chain_position) | ||
| fails = check_multiple_attrs({"posterior": ["mu", "tau", "theta"]}, idata) |
There was a problem hiding this comment.
| fails = check_multiple_attrs({"posterior": ["mu", "tau", "theta"]}, idata) | |
| fails = check_multiple_attrs({"posterior": ["mu", "tau", "theta"], "~sample_stats": []}, idata) |
check_multiple_attrs can be used to both check something is present or something is not. Here we want to make sure sample_stats is not present because there is no info argument. There is no point in having test_from_blackjax_no_info as a different test, it is doing exactly the same as this one. This extends to several other redundant tests that stem from suboptimal use of check_multiple_attrs
| if arr.ndim == 0: | ||
| result[name] = arr.reshape(1, 1) | ||
| elif num_chains is not None and arr.ndim >= 2 and arr.shape[0] == num_chains: | ||
| result[name] = arr | ||
| else: | ||
| result[name] = np.expand_dims(arr, axis=0) |
There was a problem hiding this comment.
This will end up in expand_dims whenever num_chains is not provided. I don't know much about blackjax and know it is not a PPL but a lower level library focused on sampling. However, sampling a single chain is not recommended so I would switch the API (both here and in from_blackjax) to take sample_dims as an input instead of using rc_context to ensure it is always chain and draw like the other samplers. That means by default chain and draw are assumed to be present unless the rcParam is modified or the argument passed explicitly to for example use sample_dims="sample" in which case we wouldn't need to expand the dimensions at all.
IIRC, only the loo related functionality has chain, draw dimensions as a hardcoded requirement in the input so maybe we could add an example around that in the conversion notebook. Something like 1) use sample_dims="draw", 2) use map_over_datasets to call expand_dims(chain=1) on all relevant groups.
| ------- | ||
| DataTree | ||
|
|
||
| Examples |
There was a problem hiding this comment.
can you move this to a notebook with the multiple conversion examples like the other converter functions have?
There are several things I am not completely sure about API wise and having proper examples that run as is would be helpful.
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
|
hey @OriolAbril i have addressed the review comments...
let me know if anything needs further changes |
|
5 remaining failures are all pre-existing in |
OriolAbril
left a comment
OriolAbril
left a comment
There was a problem hiding this comment.
only one high level comment regarding the API, I will review in depth after that.
Also do note this is not true:
5 remaining failures are all pre-existing in test_numpyro.py::test_mcmc_with_thinning and are unrelated to this PR. They fail because the CI machine only has 1 JAX device, but the test uses nchains=2 with chain_method="parallel", which triggers a UserWarning that gets treated as an error by pytest. This was already failing on main before my changes. All 18 blackjax tests pass.
CI on main is all green. This probably has to do with the introduction of blackjax to the env modifying the jax version and which trigers the warning. It might have happened without this PR later on but this is very different from what the message says. If you can please look into this better to see if it can be fixed robustly otherwise you'll have to wait until we can look into it. While not directly related we can't merge this if it is breaking the numpyro tests.
| num_chains : int, optional | ||
| Number of chains. Required when ``sample_dims`` includes ``"chain"`` | ||
| and the position arrays already carry a leading chain axis (i.e. | ||
| multi-chain runs via ``jax.vmap`` / ``jax.pmap``). Ignored when | ||
| ``sample_dims=["sample"]``. | ||
| max_tree_depth : int, optional | ||
| Maximum tree depth used during NUTS sampling. When provided, | ||
| ``reached_max_tree_depth`` is added to ``sample_stats``. | ||
| sample_dims : list of str, optional | ||
| Names of the sample dimensions. Defaults to | ||
| ``rcParams["data.sample_dims"]`` (typically ``["chain", "draw"]``). | ||
|
|
||
| Use ``["chain", "draw"]`` for standard multi/single-chain NUTS runs | ||
| where you want ArviZ to track the chain dimension explicitly. | ||
|
|
||
| Use ``["sample"]`` when you want a flat sample representation without | ||
| a chain dimension (e.g. after thinning or when compatibility with | ||
| LOO tools is not required). |
There was a problem hiding this comment.
This has nothing to do with #174 (comment). Take for example this piece of the comment:
That means by default chain and draw are assumed to be present unless the rcParam is modified or the argument passed explicitly
The code is not assuming chain and draw are present. It still does exactly the same as before. The default behaviour continues to be that if n_chains is not passed then the dimension is added. IIUC, taking the variables from the notebook, from_blackjax(states.position) will work but from_blackjax(mc_states.position) will not. As I said, running a single chain is not recommended, so I don't think the from_blackjax API should "reward" that by making it simpler while making the recommended behaviour harder. Given the default is sample_dims being chain and draw, the second example should work out of the box, and the first example should require being explicit about the lack of chain dimension, e.g. from_blackjax(states.position, sample_dims="draw").
Not everything we'll say will be clear but reviews allow for further comments so you can ask for clarification.
There was a problem hiding this comment.
thanks for the explanation
just to confirm: when sample_dims=["chain", "draw"] (the default) and num_chains is provided, the leading axis is treated as the chain axis. But when sample_dims=["chain", "draw"] and num_chains is NOT provided, should we raise an error asking the user to either pass num_chains or use sample_dims=["draw"]? Or should we still try to handle it somehow?
There was a problem hiding this comment.
I think there should be no num_chains argument, I don't see what value it would add if we already have sample_dims. If sample_dims has chain as the first dimension then the first dimension in the data should be assumed to be the chain dimension. If there is no chain dimension then the converter should error out. I think the behaviour on this should basically be like from_dict
There was a problem hiding this comment.
got it, i will remove num_chains entirely and update _ensure_chain_dim to simply assume the first axis is the chain dimension when sample_dims starts with "chain", erroring out if the data doesn't have enough dimensions. I'll also look into the numpyro test failures caused by the jax version change from adding blackjax to the CI environment.
2 participants
Closes #165
what this PR does
adds a
from_blackjax()converter that turns BlackJAX sampling output into an ArviZDataTree(InferenceData) object, similar to the existingfrom_numpyroconvertergroups supported
posteriorfromstates.position(dict, NamedTuple, or bare array)sample_statsfrom the BlackJAX info object (NUTS fields renamed to ArviZ standard names)prior/prior_predictivesplit by whether the variable appears in the posteriorposterior_predictive,observed_data,constant_datadesign decisions
(n_draws, *event_shape)get a chain dim of 1 inserted automatically(n_chains, n_draws, *event_shape)are detected vianum_chainsjax.device_getstates.positionandinfosdirectlylog_likelihoodis not included...BlackJAX has no model trace to compute it from; open to discussion on whether to accept it as a plain dict in a follow upnotes
from_numpyroandfrom_dictbefore writing thistests/test_io_blackjax.py