lnccbrown / lnccbrown/HSSM

MCMC sampling hangs in tests/test_save_load.py::test_save_load_vi_mcmc possibly after external dependency updates

Open
#894 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
124
Forks
24
Avg merge
19h 32m
Merged PRs (30d)
60

Description

Recently the test tests/test_save_load.py::test_save_load_vi_mcmc started hanging during PyMC’s NUTS initialization and sampling. The model and ONNX/JAX initialization succeed, but sampling does not progress (even with single-threaded execution).

Observed logs:

Model initialized successfully
ONNX/JAX backend fallback messages (TPU/GPU not available, falling back to CPU)
PyMC: “Initializing NUTS using adapt_diag...”
Multiprocess sampling (2 chains in 2 jobs)
Progress bar stuck at 0 draws

What changed:

This test previously passed; the issue seems to have appeared after updating dependencies (PyMC, JAX, ONNX, etc.).

Run pytest tests/test_save_load.py::test_save_load_vi_mcmc to reproduce. Observe hang during MCMC sampling.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running pytest tests/test_save_load.py::test_save_load_vi_mcmc and inspect the hang during PyMC NUTS initialization and sampling. Compare the dependency updates with the previously passing GitHub Actions run, focusing on the PyMC, JAX, and ONNX-related environment described in the issue. Done means the test completes successfully without hanging.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
testing
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.