Make zarr traces equivalent to arrow traces

[velochy]

What / why

A trace returned from a zarr_store (added in #244, following the request in #171) currently
diverges from the arrow backend in four ways that break downstream consumers:

1. Missing sample_stats attrs. The arviz-convention attrs added in 7451919
   (inference_library, inference_library_version, inference_library_settings) are only
   set on the arrow path. PyMC's patch_nutpie_idata reads them, so

   pm.sample(nuts_sampler="nutpie", nuts_sampler_kwargs={"zarr_store": store})

   fails with KeyError: 'inference_library_settings' — the zarr backend is unusable through
   PyMC without monkeypatching.

2. No chain/draw coordinate variables. They exist as dimensions only, so label-based
   selection (posterior.sel(draw=...)) raises KeyError on zarr traces where it works on
   arrow ones.

3. Dict-valued attrs. Fine in zarr, but with no HDF5 equivalent, so a later
   to_netcdf() of the tree fails with TypeError: Object dtype has no native HDF5 equivalent.

4. save_warmup and the store_* settings are ignored. The trace keeps both warmup groups
   and ten per-draw, parameter-sized stats (gradient, unconstrained_draw, mass_matrix_*,
   divergence_*, transformed_*) that the arrow backend drops. On a small PyMC model that is
   an 8x larger stored trace, and the ratio grows with the parameter count — it inflates every
   saved trace, every re-read, and the peak memory of writing one out.

The fix handles all four in the zarr branch of _extract. The attrs and the skip-list are hoisted
into _sample_stats_attrs() / _skipped_stats() helpers used by both paths, so the two backends
cannot drift again; chain/draw are labelled 0..n-1 as on the arrow path; dict attrs are
JSON-encoded the way inference_library_settings already is; and the stats and warmup groups the
settings exclude are dropped. Coords and vars are assigned per node in place, so the tree stays
lazy and keeps any children (tree[path] = dataset would replace the node and discard them).

Motivation/context: #171 asked for the zarr backend precisely so that "pymc would then be
able to simply load the trace from the underlying storage without much need to nutpie
wrappers" — this makes that hold in practice. It also matters for the memory threads (#233,
#265): streaming to a zarr_store keeps sampling-phase memory flat instead of accumulating
the trace in RAM, but only if the resulting trace is actually consumable downstream.

Test plan

• tests/test_pymc.py::test_zarr_store_sample_stats_attrs (new) — asserts attr parity
  between the zarr and arrow trace, that num_tune (the field PyMC dereferences)
  round-trips, that sel(draw=...) works, that the tree survives to_netcdf(), and that both
  backends return the same groups and the same sampler stats.
• Existing zarr test passes unchanged; verified end to end that
  pm.sample(..., zarr_store=...) completes with no shims on this branch.

🤖 Generated with Claude Code

[velochy]

Real Margus again:
This is a smaller and safer PR, but makes zarr much more usable with pymc - it can currently still be made to work, but some shimmying is required, whereas this makes it a drop-in replacement for most everything downstream.
Seems like an easy and safe win :)

[velochy]

Extended: found a fourth divergence while using this downstream — the zarr path also ignores save_warmup and the store_* settings, so a stored trace keeps both warmup groups and ten per-draw, parameter-sized stats the arrow backend drops (8x larger trace on a small model, worse as parameters grow). The skip-list is now a shared helper both backends call, so they can't drift apart again, and the test asserts group/stat parity. Also switched the coord labelling to assign per node in place — tree[path] = dataset replaces the node and silently drops its children.

[aseyboldt]

Thanks, this look useful. I'm wondering though if we should move most of this to the nuts-rs repo, so that users can just load the zarr file manually later if they want, and still get the correct format.
The store_* entries should be empty (ie, only nan) in the zarr case, and shouldn't take any disk or memory space. If that's not true for some reason, we need to fix it.

[velochy]

You're right — correcting my description. Six of the ten are zero-length and cost nothing; four (gradient, transformed_gradient, transformed_position, unconstrained_draw) are allocated full-length and NaN-filled, compressing to ~1 KB each in the store. They're free in zarr, but xarray materializes NaN at full width on load and on to_netcdf. On a 40-param/200-draw/2-chain fit with save_warmup=False: netCDF from arrow 0.21 MB, from zarr 1.44 MB — ~0.72 MB of that is the warmup groups (real data, kept anyway), ~0.51 MB those four arrays.

So the store_* half is a shape bug, not a volume one, and fixing it in nuts-rs (write them zero-length like the other six) makes them free everywhere rather than only in zarr — which supports moving this there, along with the chain/draw coords. The arviz sample_stats attrs and JSON-encoding dict attrs for to_netcdf are the parts that look nutpie-Python to me, but nuts-rs could write those too.

[velochy]

Took a look at nuts-rs — you're right that it belongs there, and the store_* half turned out to be a small change.

The mechanism: create_arrays makes an array per declared stat, and a disabled stat is None at extraction, so its array just stays at the NaN fill. Six of the ten are harmless — their primary dim is divergence-count, so they're already zero-length. The four draw-dimensioned ones (gradient, unconstrained_draw, transformed_position, transformed_gradient) allocate in full. As you say they're ~free in the store (all-NaN compresses to ~1 KB each), but anything reading the trace materializes them: on a 40-parameter, 200-draw, 2-chain fit that is 0.5 MB of a 1.4 MB netCDF conversion (arrow: 0.21 MB). So I'll correct my earlier claim — it's a shape issue, not data volume.

The natural hook is Settings::stat_names, where &self (carrying the store_* flags) already meets the name list create_arrays consumes. Branch here, if useful: https://github.com/velochy/nuts-rs/tree/zarr-honour-settings — it adds a disabled_stats() trait method (defaults to empty, so implementors outside the crate are unaffected), filters stat_names through it, and overrides it in the six concrete settings impls from a shared mapping mirroring TransformedPoint::extract_stats. Since stat_types/stat_dims_all all derive from stat_names, the arrow backend inherits it too and the two paths can't drift. cargo check clean, 41/41 unit tests pass — but compile-and-unit-tested only; I have not rebuilt nutpie's extension against it, so no end-to-end check yet.

The warmup half I left alone: save_warmup doesn't appear anywhere in nuts-rs, so the zarr writer creates warmup_posterior/warmup_sample_stats unconditionally. Plumbing that flag into the store config is an API call I'd rather leave to you — and it's the bigger half (0.72 of the 1.23 MB excess above).

Happy to open that as a nuts-rs PR whenever you want it, and to trim this one down to just the attrs + JSON-encoding once the Rust side lands.