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

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions