From 14a35b36b7a9244372bef29056df9b060a1e9b14 Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Tue, 28 Jul 2026 18:40:51 +0300 Subject: [PATCH 1/4] fix: make zarr traces equivalent to arrow traces A trace returned from a zarr_store diverged from the arrow backend in four ways that break downstream consumers: * the arviz-convention sample_stats attrs added in 74519195 were only set on the arrow path, so pymc's patch_nutpie_idata fails with KeyError: 'inference_library_settings' - the zarr backend is unusable through pm.sample without monkeypatching; * chain/draw exist as dimensions but carry no coordinate variables, so label-based selection (posterior.sel(draw=...)) raises KeyError; * dict-valued attrs serialize fine to zarr but have no HDF5 equivalent, so a later to_netcdf() of the tree fails with 'Object dtype has no native HDF5 equivalent'; * save_warmup and the store_* settings are ignored, so the trace keeps both warmup groups and ten per-draw, parameter-sized sampler stats (gradient, unconstrained_draw, mass_matrix_*, divergence_*, transformed_*) that the arrow backend drops. Measured on a small pymc model that is an 8x larger stored trace, and the ratio grows with the parameter count. Hoist the attrs and the skip-list into helpers used by both paths, label chain/draw as the arrow backend does, JSON-encode dict attrs the way inference_library_settings already is, and drop the stats and warmup groups the settings exclude. Coords and vars are assigned per node in place, so the tree stays lazy and keeps any children. The regression test asserts attr parity with arrow, that num_tune (the field pymc dereferences) round-trips, that label-based selection works, that the tree survives to_netcdf, and that both backends return the same groups and sampler stats. Co-Authored-By: Claude Fable 5 --- python/nutpie/sample.py | 120 ++++++++++++++++++++++++++-------------- tests/test_pymc.py | 50 +++++++++++++++++ 2 files changed, 129 insertions(+), 41 deletions(-) diff --git a/python/nutpie/sample.py b/python/nutpie/sample.py index 10b5c25..a94010c 100644 --- a/python/nutpie/sample.py +++ b/python/nutpie/sample.py @@ -605,6 +605,51 @@ def wait(self, *, timeout=None): results = self._sampler.take_results() return self._extract(results) + def _sample_stats_attrs(self): + from nutpie import __version__ + + return { + "inference_library": "nutpie", + "inference_library_version": __version__, + "inference_library_settings": json.dumps(self._settings.as_dict()), + } + + def _skipped_stats(self, settings_dict): + """Sampler stats the settings say not to store, and so should not reach the trace.""" + skips = { + "store_gradient": ["gradient"], + "store_unconstrained": ["unconstrained_draw"], + "adapt_options.mass_matrix_options.store_mass_matrix": [ + "mass_matrix_inv", + "mass_matrix_eigvals", + "mass_matrix_stds", + ], + "store_divergences": [ + "divergence_start", + "divergence_end", + "divergence_momentum", + "divergence_start_gradient", + ], + "store_transformed": [ + "transformed_position", + "transformed_gradient", + "transformation_mu", + ], + } + + def _get_nested(settings, name, default): + for part in name.split("."): + if part not in settings: + return default + settings = settings[part] + return settings + + skip_vars = [] + for setting, names in skips.items(): + if not _get_nested(settings_dict["settings"], setting, False): + skip_vars.extend(names) + return skip_vars + def _extract(self, results): settings_dict = self._settings.as_dict() if self._return_raw_trace: @@ -622,52 +667,45 @@ def _extract(self, results): store = cls(*args, **kwargs) obj_store = ObjectStore(store, read_only=True) - return xr.open_datatree(obj_store, engine="zarr", consolidated=False) # ty:ignore[invalid-argument-type] + trace = xr.open_datatree(obj_store, engine="zarr", consolidated=False) # ty:ignore[invalid-argument-type] + # match the arrow backend, pymc reads these from sample_stats + if "sample_stats" in trace: + trace["sample_stats"].attrs.update(self._sample_stats_attrs()) + # the settings say which stats to store; the zarr writer stores them all + skip_vars = set(self._skipped_stats(settings_dict)) + # label chain/draw as the arrow backend does (assigning in place stays lazy and + # keeps each node's children) + for node in trace.subtree: + dataset = node.dataset + missing = { + dim: np.arange(dataset.sizes[dim]) + for dim in ("chain", "draw") + if dim in dataset.dims and dim not in dataset.coords + } + stale = ( + skip_vars & set(dataset.data_vars) + if (node.name or "").endswith("sample_stats") + else set() + ) + if missing or stale: + node.dataset = dataset.drop_vars(stale).assign_coords(missing) + if not self._save_warmup: + for name in ("warmup_posterior", "warmup_sample_stats"): + if name in trace: + del trace[name] + # dict attrs have no HDF5 equivalent; keep the tree to_netcdf-able + for node in trace.subtree: + for key, value in list(node.attrs.items()): + if isinstance(value, dict): + node.attrs[key] = json.dumps(value) + return trace elif results.is_arrow(): - skip_vars = [] - skips = { - "store_gradient": ["gradient"], - "store_unconstrained": ["unconstrained_draw"], - "adapt_options.mass_matrix_options.store_mass_matrix": [ - "mass_matrix_inv", - "mass_matrix_eigvals", - "mass_matrix_stds", - ], - "store_divergences": [ - "divergence_start", - "divergence_end", - "divergence_momentum", - "divergence_start_gradient", - ], - "store_transformed": [ - "transformed_position", - "transformed_gradient", - "transformation_mu", - ], - } - - def _get_nested(settings, name, default): - parts = name.split(".") - for part in parts: - if part not in settings: - return default - settings = settings[part] - return settings - - for setting, names in skips.items(): - if not _get_nested(settings_dict["settings"], setting, False): - skip_vars.extend(names) + skip_vars = self._skipped_stats(settings_dict) draw_batches, stat_batches = results.get_arrow_trace() - from nutpie import __version__ - - attrs = { - "inference_library": "nutpie", - "inference_library_version": __version__, - "inference_library_settings": json.dumps(self._settings.as_dict()), - } + attrs = self._sample_stats_attrs() return _arrow_to_arviz( draw_batches, diff --git a/tests/test_pymc.py b/tests/test_pymc.py index 9a087d8..3a331ba 100644 --- a/tests/test_pymc.py +++ b/tests/test_pymc.py @@ -1,3 +1,4 @@ +import json import time from importlib.util import find_spec @@ -650,3 +651,52 @@ def test_unnamed_shared(backend, gradient_backend): compiled = nutpie.compile_pymc_model(model) nutpie.sample(compiled) + + +def test_zarr_store_sample_stats_attrs(tmp_path): + """The zarr backend must attach the same sample_stats attrs as the arrow backend. + + pymc's ``patch_nutpie_idata`` reads ``inference_library_settings`` from there, so + ``pm.sample(nuts_sampler="nutpie", nuts_sampler_kwargs={"zarr_store": ...})`` + fails with a KeyError when they are missing. + """ + with pm.Model() as model: + pm.Normal("x") + + compiled = nutpie.compile_pymc_model(model, backend="numba") + + path = tmp_path / "trace.zarr" + path.mkdir() + store = nutpie.zarr_store.LocalStore(str(path)) + zarr_trace = nutpie.sample( + compiled, chains=1, seed=123, draws=20, tune=20, zarr_store=store + ) + arrow_trace = nutpie.sample(compiled, chains=1, seed=123, draws=20, tune=20) + + for key in [ + "inference_library", + "inference_library_version", + "inference_library_settings", + ]: + assert key in zarr_trace.sample_stats.attrs, key + assert ( + zarr_trace.sample_stats.attrs[key] == arrow_trace.sample_stats.attrs[key] + ), key + + settings = json.loads(zarr_trace.sample_stats.attrs["inference_library_settings"]) + assert settings["settings"]["num_tune"] == 20 + + # the store keeps every stat and both warmup groups; the settings say otherwise + assert set(zarr_trace.sample_stats.data_vars) == set(arrow_trace.sample_stats.data_vars) + assert set(zarr_trace.children) == set(arrow_trace.children) + + # chain/draw carry coordinate variables, so label-based selection works + assert "chain" in zarr_trace.posterior.coords + assert "draw" in zarr_trace.posterior.coords + zarr_trace.posterior.sel(draw=5) + + # no dict-valued attrs remain anywhere: the tree must survive to_netcdf + for node in zarr_trace.subtree: + for value in node.attrs.values(): + assert not isinstance(value, dict) + zarr_trace.to_netcdf(tmp_path / "roundtrip.nc") From 0a0da3a6916d5dfc15591dccfe875e76bb2f0919 Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Fri, 7 Aug 2026 11:20:21 +0300 Subject: [PATCH 2/4] fix: move unconstrained value variables out of the zarr posterior The zarr writer stores every variable it is handed, including each free RV's transformed value var. The arrow path pops those out via reparameterized_names and only resurfaces them, as unconstrained_posterior, when store_unconstrained is set; the zarr path returned the tree verbatim, so store_unconstrained meant different things on the two backends. Use the same reparameterized_names list rather than a name-suffix heuristic, so the two paths cannot disagree about what counts as unconstrained. --- python/nutpie/sample.py | 16 ++++++++++++++ tests/test_pymc.py | 46 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/python/nutpie/sample.py b/python/nutpie/sample.py index a94010c..b8e501c 100644 --- a/python/nutpie/sample.py +++ b/python/nutpie/sample.py @@ -693,6 +693,22 @@ def _extract(self, results): for name in ("warmup_posterior", "warmup_sample_stats"): if name in trace: del trace[name] + # the zarr writer leaves the unconstrained value variables in `posterior`; + # the arrow backend moves them out, into their own group when asked for them + uc_names = self._compiled_model.reparameterized_names or [] + for group, uc_group in ( + ("posterior", "unconstrained_posterior"), + ("warmup_posterior", "warmup_unconstrained_posterior"), + ): + if group not in trace: + continue + dataset = trace[group].dataset + present = [name for name in uc_names if name in dataset.data_vars] + if not present: + continue + if self._store_unconstrained: + trace[uc_group] = xr.DataTree(dataset[present]) + trace[group].dataset = dataset.drop_vars(present) # dict attrs have no HDF5 equivalent; keep the tree to_netcdf-able for node in trace.subtree: for key, value in list(node.attrs.items()): diff --git a/tests/test_pymc.py b/tests/test_pymc.py index 3a331ba..bd8bb01 100644 --- a/tests/test_pymc.py +++ b/tests/test_pymc.py @@ -687,7 +687,9 @@ def test_zarr_store_sample_stats_attrs(tmp_path): assert settings["settings"]["num_tune"] == 20 # the store keeps every stat and both warmup groups; the settings say otherwise - assert set(zarr_trace.sample_stats.data_vars) == set(arrow_trace.sample_stats.data_vars) + assert set(zarr_trace.sample_stats.data_vars) == set( + arrow_trace.sample_stats.data_vars + ) assert set(zarr_trace.children) == set(arrow_trace.children) # chain/draw carry coordinate variables, so label-based selection works @@ -700,3 +702,45 @@ def test_zarr_store_sample_stats_attrs(tmp_path): for value in node.attrs.values(): assert not isinstance(value, dict) zarr_trace.to_netcdf(tmp_path / "roundtrip.nc") + + +def test_zarr_store_transformed_variables(tmp_path): + """Unconstrained value variables belong out of ``posterior``, on both backends. + + The zarr writer stores every variable it is handed, including each free RV's + transformed value var; the arrow backend pops those into ``unconstrained_posterior`` + and only keeps them when ``store_unconstrained=True``. + """ + with pm.Model() as model: + sigma = pm.HalfNormal("sigma") # gives a sigma_log__ value variable + pm.Normal("mu", 0.0, sigma) + + compiled = nutpie.compile_pymc_model(model, backend="numba") + + def fit(store_unconstrained, name): + path = tmp_path / name + path.mkdir() + kwargs = dict( + chains=1, + seed=123, + draws=20, + tune=20, + store_unconstrained=store_unconstrained, + ) + store = nutpie.zarr_store.LocalStore(str(path)) + return ( + nutpie.sample(compiled, zarr_store=store, **kwargs), + nutpie.sample(compiled, **kwargs), + ) + + zarr_trace, arrow_trace = fit(False, "default.zarr") + assert set(zarr_trace.posterior.data_vars) == set(arrow_trace.posterior.data_vars) + assert "sigma_log__" not in zarr_trace.posterior.data_vars + assert "unconstrained_posterior" not in zarr_trace.children + + zarr_trace, arrow_trace = fit(True, "unconstrained.zarr") + assert "sigma_log__" not in zarr_trace.posterior.data_vars + assert "sigma_log__" in zarr_trace.unconstrained_posterior.data_vars + assert set(zarr_trace.unconstrained_posterior.data_vars) == set( + arrow_trace.unconstrained_posterior.data_vars + ) From 7233b1f0aea9a084010761adb4cc7bc996621ae8 Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Tue, 11 Aug 2026 11:03:11 +0300 Subject: [PATCH 3/4] fix: gate transformation_mu on store_mass_matrix, not store_transformed nuts-rs produces transformation_mu when store_mass_matrix is set (src/transform/diagonal.rs:57), but the skip table listed it under store_transformed. With store_mass_matrix=True and store_transformed left at its default False - the natural combination - nutpie deleted a stat the sampler had written, on both backends. Verified: present=False before, present=True after, arrow and zarr alike. Co-Authored-By: Claude Opus 5 (1M context) --- python/nutpie/sample.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/python/nutpie/sample.py b/python/nutpie/sample.py index b8e501c..2b0b766 100644 --- a/python/nutpie/sample.py +++ b/python/nutpie/sample.py @@ -623,6 +623,7 @@ def _skipped_stats(self, settings_dict): "mass_matrix_inv", "mass_matrix_eigvals", "mass_matrix_stds", + "transformation_mu", ], "store_divergences": [ "divergence_start", @@ -630,11 +631,7 @@ def _skipped_stats(self, settings_dict): "divergence_momentum", "divergence_start_gradient", ], - "store_transformed": [ - "transformed_position", - "transformed_gradient", - "transformation_mu", - ], + "store_transformed": ["transformed_position", "transformed_gradient"], } def _get_nested(settings, name, default): From 9b059c0fe0ebbe63cf3359ac08737873fde5eefc Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Tue, 11 Aug 2026 11:03:11 +0300 Subject: [PATCH 4/4] test: mark the new zarr tests so CI runs them Both tests lacked @pytest.mark.pymc, and CI only ever runs marker-selected subsets (pytest -m "pymc and not flow"), so neither ran - the whole test plan for this branch was deselected. Also fix the ruff C408 failure that blocks the lint hook, and trim the two docstrings to the 1-2 lines the repo uses. Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_pymc.py | 32 +++++++++++++------------------- 1 file changed, 13 insertions(+), 19 deletions(-) diff --git a/tests/test_pymc.py b/tests/test_pymc.py index bd8bb01..e033a23 100644 --- a/tests/test_pymc.py +++ b/tests/test_pymc.py @@ -653,13 +653,10 @@ def test_unnamed_shared(backend, gradient_backend): nutpie.sample(compiled) +@pytest.mark.pymc def test_zarr_store_sample_stats_attrs(tmp_path): - """The zarr backend must attach the same sample_stats attrs as the arrow backend. - - pymc's ``patch_nutpie_idata`` reads ``inference_library_settings`` from there, so - ``pm.sample(nuts_sampler="nutpie", nuts_sampler_kwargs={"zarr_store": ...})`` - fails with a KeyError when they are missing. - """ + """The zarr trace must carry the same sample_stats attrs, coords and stat set as arrow: + pymc's ``patch_nutpie_idata`` reads ``inference_library_settings`` from there.""" with pm.Model() as model: pm.Normal("x") @@ -704,13 +701,10 @@ def test_zarr_store_sample_stats_attrs(tmp_path): zarr_trace.to_netcdf(tmp_path / "roundtrip.nc") +@pytest.mark.pymc def test_zarr_store_transformed_variables(tmp_path): - """Unconstrained value variables belong out of ``posterior``, on both backends. - - The zarr writer stores every variable it is handed, including each free RV's - transformed value var; the arrow backend pops those into ``unconstrained_posterior`` - and only keeps them when ``store_unconstrained=True``. - """ + """Unconstrained value variables belong out of ``posterior`` on both backends, and in + ``unconstrained_posterior`` only when ``store_unconstrained=True``.""" with pm.Model() as model: sigma = pm.HalfNormal("sigma") # gives a sigma_log__ value variable pm.Normal("mu", 0.0, sigma) @@ -720,13 +714,13 @@ def test_zarr_store_transformed_variables(tmp_path): def fit(store_unconstrained, name): path = tmp_path / name path.mkdir() - kwargs = dict( - chains=1, - seed=123, - draws=20, - tune=20, - store_unconstrained=store_unconstrained, - ) + kwargs = { + "chains": 1, + "seed": 123, + "draws": 20, + "tune": 20, + "store_unconstrained": store_unconstrained, + } store = nutpie.zarr_store.LocalStore(str(path)) return ( nutpie.sample(compiled, zarr_store=store, **kwargs),