Skip to content

diagnostics.to_datatree only accepts MCMCResult — no ArviZ export path for SVI/Predictive fits #43

Description

@aaronstevenwhite

Summary

quivers.diagnostics.to_datatree accepts only an MCMCResult, so there is no supported path to export an SVI fit (an inference.guides Guide / Predictive draws) to ArviZ, even though the module docstring notes the detach in _tensor_to_numpy is "necessary because samples may carry an autograd graph (e.g. when produced by SVI's Predictive rather than NUTS)."

Evidence

src/quivers/diagnostics/arviz_io.py:

def to_datatree(posterior: MCMCResult, *, ...):
    ...
    for name, t in posterior.samples.items(): ...
    sample_stats_group = {"lp": _tensor_to_numpy(posterior.log_densities), ...}
    ... posterior.acceptance_rates ... posterior.divergence_counts ...

It hard-requires MCMCResult-only attributes (.samples, .log_densities, .acceptance_rates, .divergence_counts, .num_chains, .num_samples). A Guide has none of these, so SVI users must hand-roll az.from_dict.

Repro

After an SVI fit, Predictive(model, guide, num_samples=...)(x, obs) yields a dict[str, Tensor] of draws, but to_datatree(...) cannot consume it — there is no posterior: Guide | Predictive overload.

Expected

Either an overload / second entry point that builds a DataTree from variational draws (a guide + Predictive output, with a single pseudo-chain), or a documented recipe. This matters because _fit_svi is the default for quivers.formulas.fit(..., method="svi"), so the most common fit path has no first-class ArviZ export.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions