Migrate to support Arviz 1.0

[zaxtax]

Arviz is preparing for a backwards-compatibility breaking 1.0 release. We should migrate pymc-extras to support it.

The migration guide seems pretty decent as a place to start for that. Though, I suspect we do this piecemeal between the example notebooks and the codebase itself.

[AAlexxis222]

Hi @zaxtax! I would like to help with this.

I have just finished refactoring 34 plotting functions in arviz-plots to enforce the new signature standards (keyword-only arguments, sample_dims positioning, etc.) for the ArviZ 1.0 roadmap (PR merged yesterday).

Because of this, I am very familiar with the recent breaking changes and the new API patterns. I can start by auditing the pymc-extras codebase to identify where these plot functions are used and apply the necessary updates.

Please let me know if I can go ahead!

[aloctavodia]

@AAlexxis222 notice that most of the needed changes are most likely not related to plotting functions, but to the adoption of xarray's datatree in replacement of InferenceData. You may want to evaluate which tests fail when using arviz 1.0.rc0.

[AAlexxis222]

Report on ArviZ 1.0rc0 compatibility check:

I set up a fresh environment with python 3.13, pymc 5.27.1, and arviz 1.0.0rc0 as requested.

However, I cannot run the pymc-extras test suite because PyMC itself fails to import with this ArviZ version. It raises an AttributeError during the collection phase, originating from pymc core, not pymc-extras:

AttributeError: module 'arviz' has no attribute 'InferenceData'

It seems pymc-extras cannot be migrated/tested until pymc is updated to support ArviZ 1.0.

Question: Is there a specific branch of pymc (dev/experimental) that already supports ArviZ 1.0 that I should install instead of the stable release? Or should we wait for a PyMC update first?

[aloctavodia]

We have this pymc-devs/pymc#8019

[AAlexxis222]

Using PyMC PR #8019 resolved the blocking AttributeError.

Now the test suite fails because InferenceData has moved to arviz_base. The affected files are:

• pymc_extras/statespace/core/statespace.py
• pymc_extras/model/marginal/marginal_model.py
• pymc_extras/inference/dadvi/dadvi.py

Should I implement a backward-compatible shim (try/except import) to keep supporting older ArviZ versions, or should I update the imports directly (dropping support for ArviZ < 1.0)?

[aloctavodia]

InferenceData no longer exists, we now use xarray's DataTree.

[AAlexxis222]

Understood. I have refactored the codebase to use xarray.DataTree as requested (replacing imports and updating attribute access to dictionary style).

However, I encountered a blocking issue in the PyMC PR #8019 branch:
The tests fail with TypeError: _sample_return() got an unexpected keyword argument 'quiet'. It seems pm.sample() calls _sample_return(..., quiet=quiet), but the definition no longer accepts that argument.

I have patched it locally to proceed with the remaining pymc-extras fixes.

[aloctavodia]

Probably that PR needs to be rebased

[AAlexxis222]

Thanks! I have completed the migration to DataTree (imports, attributes, and constructors).

With my local patch, ~318 tests pass. However, CI will fail without the upstream fix.

Should I open the PR now so you can review the code, or wait for the PyMC rebase?

[aloctavodia]

I did the rebase.

[AAlexxis222]

PR opened: #642

Changes are live and ready for review.
(Note: I see the pre-commit failure, working on fixing the formatting now).