Skip to content

Add more idata attributes for JAX samplers #7360

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Jun 17, 2024
Merged

Add more idata attributes for JAX samplers #7360

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
a9bfe9f
fix: Include `posterior_attrs` in trace dict
osyuksel Jun 15, 2024
bb3a9c2
test: Compare external samplers to pymc reference
osyuksel Jun 16, 2024
c8ccea4
test: Check tuning steps
osyuksel Jun 16, 2024
c41ee0f
fix: Add tuning_steps for nutpie
osyuksel Jun 16, 2024
949455b
fix: Add tuning_steps for jax
osyuksel Jun 16, 2024
bbe0fe9
style: Add trailing comma to JAX attrs
osyuksel Jun 16, 2024
b8e157d
style: Add trailing comma to MCMC attrs
osyuksel Jun 16, 2024
fda9222
style: Multi-line formatting for JAX attrs
osyuksel Jun 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions pymc/sampling/jax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -669,9 +669,7 @@ def sample_jax_nuts(
else:
log_likelihood = None

attrs = {
"sampling_time": (tic2 - tic1).total_seconds(),
}
attrs = {"sampling_time": (tic2 - tic1).total_seconds(), "tuning_steps": tune}

coords, dims = coords_and_dims_for_inferencedata(model)
# Update 'coords' and 'dims' extracted from the model with user 'idata_kwargs'
Expand All @@ -680,6 +678,7 @@ def sample_jax_nuts(
coords.update(idata_kwargs.pop("coords"))
if "dims" in idata_kwargs:
dims.update(idata_kwargs.pop("dims"))

# Use 'partial' to set default arguments before passing 'idata_kwargs'
to_trace = partial(
az.from_dict,
Expand Down
4 changes: 1 addition & 3 deletions pymc/sampling/mcmc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -334,9 +334,7 @@ def _sample_external_nuts(
default_dims=[],
)
attrs = make_attrs(
{
"sampling_time": t_sample,
},
{"sampling_time": t_sample, "tuning_steps": tune},
library=nutpie,
)
for k, v in attrs.items():
Expand Down
7 changes: 7 additions & 0 deletions tests/sampling/test_mcmc_external.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,10 @@ def test_external_nuts_sampler(recwarn, nuts_sampler):
idata1 = sample(**kwargs)
idata2 = sample(**kwargs)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not introduced by you, but why are we sampling twice, is idata2 used anywhere?


reference_kwargs = kwargs.copy()
reference_kwargs["nuts_sampler"] = "pymc"
idata_reference = sample(**reference_kwargs)

warns = {
(warn.category, warn.message.args[0])
for warn in recwarn
Expand All @@ -64,8 +68,11 @@ def test_external_nuts_sampler(recwarn, nuts_sampler):
assert "L" in idata1.observed_data
assert idata1.posterior.chain.size == 2
assert idata1.posterior.draw.size == 500
assert idata1.posterior.tuning_steps == 500
np.testing.assert_array_equal(idata1.posterior.x, idata2.posterior.x)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ricardoV94 idata2 is used in this line only, to check if the sampling is deterministic I think?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah okay! Thanks for checking


assert idata_reference.posterior.attrs.keys() == idata1.posterior.attrs.keys()


def test_step_args():
with Model() as model:
Expand Down