Add more idata attributes for JAX samplers (#7360)

Co-authored-by: Ricardo Vieira <[email protected]>

Author: osyuksel

pymc/sampling/jax.py

@@ -671,6 +671,7 @@ def sample_jax_nuts(
 
     attrs = {
         "sampling_time": (tic2 - tic1).total_seconds(),
+        "tuning_steps": tune,
     }
 
     coords, dims = coords_and_dims_for_inferencedata(model)
@@ -680,6 +681,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,
@@ -690,6 +692,7 @@ def sample_jax_nuts(
         coords=coords,
         dims=dims,
         attrs=make_attrs(attrs, library=library),
+        posterior_attrs=make_attrs(attrs, library=library),
     )
     az_trace = to_trace(posterior=mcmc_samples, **idata_kwargs)
 

pymc/sampling/mcmc.py

@@ -336,6 +336,7 @@ def _sample_external_nuts(
         attrs = make_attrs(
             {
                 "sampling_time": t_sample,
+                "tuning_steps": tune,
             },
             library=nutpie,
         )

tests/sampling/test_mcmc_external.py

@@ -44,6 +44,10 @@ def test_external_nuts_sampler(recwarn, nuts_sampler):
         idata1 = sample(**kwargs)
         idata2 = sample(**kwargs)
 
+        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
@@ -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)
 
+    assert idata_reference.posterior.attrs.keys() == idata1.posterior.attrs.keys()
+
 
 def test_step_args():
     with Model() as model: