Initial migration to pymc6

[digicosmos86]

This is the first of many PRs required to get HSSM to be compatible with PyMC6. All future PRs need to be merged to migration-pymc6 branch

Changes in this PR:

1. updated dependencies in pyproject.toml (still many things need to be decided, such as whether to drop numpyro in favor of nutpie
2. Updated a few utilities that post-process InferenceData object. Now that the InferenceData is replace with a xr.DataTree object, most functions that deals with the InferenceData object need to be updated
3. I dug a little bit into the main tutorial. Many arviz functions need to be updated

[Copilot]

Pull request overview

This PR starts the migration of HSSM toward compatibility with the newer PyMC/Bambi/ArviZ stack by updating dependencies, switching parts of the trace/log-likelihood pipeline from ArviZ InferenceData to xarray.DataTree, and adjusting a test import for the updated PyTensor module layout.

Changes:

• Update project dependency constraints and Python minimum version in pyproject.toml.
• Refactor log-likelihood computation helper to operate on xr.DataTree instead of az.InferenceData.
• Update sampling post-processing to clean the posterior group using DataTree-style access, plus a PyTensor test import adjustment.

Reviewed changes

Copilot reviewed 4 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/test_likelihoods.py	Updates NanGuardMode import path for PyTensor’s new module location.
src/hssm/utils.py	Refactors _compute_log_likelihood to use xr.DataTree and updates log-likelihood group assignment.
src/hssm/base.py	Updates sample() return typing and switches posterior cleanup to DataTree-style manipulation.
pyproject.toml	Adjusts Python requirement and dependency set for the migration.
.gitignore	Ignores uv.lock.

Comments suppressed due to low confidence (2)

src/hssm/base.py:557

• The return-type annotation for sample() was updated to xr.DataTree | pm.Approximation, but the docstring still states that model.traces is an ArviZ InferenceData for most samplers. This is misleading for users and should be updated to match the new trace container type.

        xr.DataTree | pm.Approximation
            A reference to the `model.traces` object, which stores the traces of the
            last call to `model.sample()`. `model.traces` is an ArviZ `InferenceData`
            instance if `sampler` is `"pymc"` (default), `"numpyro"`,
            `"blackjax"` or "`laplace".

pyproject.toml:38

• pymc and arviz were removed from the project dependencies, but they are imported throughout the package (e.g., src/hssm/base.py, src/hssm/utils.py). Relying on transitive dependencies from bambi is fragile and can break installs if upstream changes extras/optional deps; declare these as direct dependencies again.

dependencies = [
    "absl-py>=2.3.1",
    "bambi>=0.18.0",
    "cloudpickle>=3.0.0",
    "hddm-wfpt>=0.1.6",
    "huggingface-hub>=1.17.0",
    "jaxonnxruntime>=0.3.0",
    "numpyro>=0.19",
    "onnx>=1.16.0",
    "pandas>=2.2,<3",
    "seaborn>=0.13.2",
    "ssm-simulators>=0.12.2",
    "tqdm>=4.66.0",
]

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[cpaniaguam]

Consider marking the yet-to-address broken tests as failing.

[digicosmos86]

Consider marking the yet-to-address broken tests as failing.

Through what mechanism? Commenting in code?

[cpaniaguam]

@digicosmos86

Through what mechanism? Commenting in code?

You can use pytest.xfail.

[cpaniaguam]

@digicosmos86

Through what mechanism? Commenting in code?

See https://docs.pytest.org/en/stable/how-to/skipping.html#xfail

[digicosmos86]

@digicosmos86

Through what mechanism? Commenting in code?

See https://docs.pytest.org/en/stable/how-to/skipping.html#xfail

I think that's fair only when there's not a whole lot of tests breaking. Otherwise adding the marks itself could be labor intensive. Why don''t we evaluate after merging your recent PRs? Looks like a few broken tests were fixed in them

[cpaniaguam]

Why don''t we evaluate after merging your recent PRs? Looks like a few broken tests were fixed in them

Those were just two of the fast tests. There are many other slow ones that are failing (see #982).

I think that's fair only when there's not a whole lot of tests breaking.

I understand your point. I still think there’s value in marking the failures we already understand, since it keeps CI easier to read and helps distinguish known issues from new regressions. Once a fix is introduced, those tests will show up as xpassed in the summary, so we’ll have a clear signal to remove the markers.