test: mark tests as pymc or stan and select in ci

Author: aseyboldt

.github/workflows/ci.yml

@@ -67,9 +67,14 @@ jobs:
           set -e
           python3 -m venv .venv
           source .venv/bin/activate
+          uv pip install 'nutpie[stan]' --find-links dist --force-reinstall
+          uv pip install pytest pytest-timeout
+          pytest -m "stan and not flow"
+          uv pip install 'nutpie[pymc]' --find-links dist --force-reinstall
+          uv pip install jax
+          pytest -m "pymc and not flow"
           uv pip install 'nutpie[all]' --find-links dist --force-reinstall
-          uv pip install pytest
-          pytest
+          pytest -m flow
       - name: pytest
         if: ${{ !startsWith(matrix.platform.target, 'x86') && matrix.platform.target != 'ppc64' }}
         uses: uraimo/run-on-arch-action@v3
@@ -85,9 +90,9 @@ jobs:
           run: |
             set -e
             source $HOME/.local/bin/env
-            uv pip install --system -U pip pytest
+            uv pip install --system -U pip pytest pytest-timeout
             uv pip install --system 'nutpie[all]' --find-links dist --force-reinstall
-            pytest -m "not slow"
+            pytest -m "not flow"  # Skip flow tests, they are slow on emulated platforms
 
   # pyarrow doesn't currently seem to work on musllinux
   #musllinux:
@@ -208,7 +213,7 @@ jobs:
           python3 -m venv .venv
           source .venv/Scripts/activate
           uv pip install "nutpie[all]" --find-links dist --force-reinstall
-          uv pip install pytest
+          uv pip install pytest pytest-timeout
           pytest
 
   macos:
@@ -252,8 +257,8 @@ jobs:
           python3 -m venv .venv
           source .venv/bin/activate
           uv pip install 'nutpie[all]' --find-links dist --force-reinstall
-          uv pip install pytest
-          pytest
+          uv pip install pytest pytest-timeout
+          pytest -m "not (flow and stan)"  # The stan tests seem to run out of memory on macOS?
 
   sdist:
     runs-on: ubuntu-latest

Cargo.lock

@@ -35,6 +35,7 @@ dev = [
   "jax >= 0.4.27",
   "flowjax >= 17.0.2",
   "pytest",
+  "pytest-timeout",
 ]
 all = [
   "bridgestan >= 2.6.1",
@@ -65,3 +66,10 @@ venv = "default"
 module-name = "nutpie._lib"
 python-source = "python"
 features = ["pyo3/extension-module"]
+
+[tool.pytest.ini_options]
+markers = [
+    "flow: tests for normalizing flows",
+    "stan: tests for Stan models",
+    "pymc: tests for PyMC models",
+]

pyproject.toml

@@ -12,7 +12,6 @@
 import numpy as np
 import pandas as pd
 from numpy.typing import NDArray
-from pymc.initial_point import make_initial_point_fn
 
 from nutpie import _lib
 from nutpie.compiled_pyfunc import SeedType, from_pyfunc
@@ -501,6 +500,7 @@ def compile_pymc_model(
         )
 
     from pymc.model.transform.optimization import freeze_dims_and_data
+    from pymc.initial_point import make_initial_point_fn
 
     if freeze_model is None:
         freeze_model = backend == "jax"

python/nutpie/compile_pymc.py

@@ -1,3 +1,10 @@
+from importlib.util import find_spec
+import time
+import pytest
+
+if find_spec("pymc") is None:
+    pytest.skip("Skip pymc tests", allow_module_level=True)
+
 import numpy as np
 import pymc as pm
 import pytest
@@ -12,6 +19,7 @@
 )
 
 
+@pytest.mark.pymc
 @parameterize_backends
 def test_pymc_model(backend, gradient_backend):
     with pm.Model() as model:
@@ -24,6 +32,7 @@ def test_pymc_model(backend, gradient_backend):
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.pymc
 @parameterize_backends
 def test_pymc_model_float32(backend, gradient_backend):
     import pytensor
@@ -39,6 +48,7 @@ def test_pymc_model_float32(backend, gradient_backend):
         trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.pymc
 @parameterize_backends
 def test_pymc_model_no_prior(backend, gradient_backend):
     with pm.Model() as model:
@@ -52,6 +62,7 @@ def test_pymc_model_no_prior(backend, gradient_backend):
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.pymc
 @parameterize_backends
 def test_blocking(backend, gradient_backend):
     with pm.Model() as model:
@@ -65,34 +76,41 @@ def test_blocking(backend, gradient_backend):
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.pymc
 @parameterize_backends
-@pytest.mark.timeout(2)
+@pytest.mark.timeout(20)
 def test_wait_timeout(backend, gradient_backend):
     with pm.Model() as model:
         pm.Normal("a", shape=100_000)
     compiled = nutpie.compile_pymc_model(
         model, backend=backend, gradient_backend=gradient_backend
     )
+    start = time.time()
     sampler = nutpie.sample(compiled, chains=1, blocking=False)
     with pytest.raises(TimeoutError):
         sampler.wait(timeout=0.1)
     sampler.cancel()
+    assert start - time.time() < 5
 
 
+@pytest.mark.pymc
 @parameterize_backends
-@pytest.mark.timeout(2)
+@pytest.mark.timeout(20)
 def test_pause(backend, gradient_backend):
     with pm.Model() as model:
         pm.Normal("a", shape=100_000)
     compiled = nutpie.compile_pymc_model(
         model, backend=backend, gradient_backend=gradient_backend
     )
+    start = time.time()
     sampler = nutpie.sample(compiled, chains=1, blocking=False)
     sampler.pause()
     sampler.resume()
     sampler.cancel()
+    assert start - time.time() < 5
 
 
+@pytest.mark.pymc
 @parameterize_backends
 def test_pymc_model_with_coordinate(backend, gradient_backend):
     with pm.Model() as model:
@@ -106,6 +124,7 @@ def test_pymc_model_with_coordinate(backend, gradient_backend):
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.pymc
 @parameterize_backends
 def test_pymc_model_store_extra(backend, gradient_backend):
     with pm.Model() as model:
@@ -130,6 +149,7 @@ def test_pymc_model_store_extra(backend, gradient_backend):
     _ = trace.sample_stats.mass_matrix_inv
 
 
+@pytest.mark.pymc
 @parameterize_backends
 def test_trafo(backend, gradient_backend):
     with pm.Model() as model:
@@ -142,6 +162,7 @@ def test_trafo(backend, gradient_backend):
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.pymc
 @parameterize_backends
 def test_det(backend, gradient_backend):
     with pm.Model() as model:
@@ -156,6 +177,7 @@ def test_det(backend, gradient_backend):
     assert trace.posterior.b.shape[-1] == 2
 
 
+@pytest.mark.pymc
 @parameterize_backends
 def test_non_identifier_names(backend, gradient_backend):
     with pm.Model() as model:
@@ -172,6 +194,7 @@ def test_non_identifier_names(backend, gradient_backend):
     assert trace.posterior["foo::b"].shape[-1] == 2
 
 
+@pytest.mark.pymc
 @parameterize_backends
 def test_pymc_model_shared(backend, gradient_backend):
     with pm.Model() as model:
@@ -197,6 +220,7 @@ def test_pymc_model_shared(backend, gradient_backend):
         nutpie.sample(compiled3, chains=1)
 
 
+@pytest.mark.pymc
 @parameterize_backends
 def test_pymc_var_names(backend, gradient_backend):
     with pm.Model() as model:
@@ -244,7 +268,8 @@ def test_pymc_var_names(backend, gradient_backend):
     assert not hasattr(trace.posterior, "c")
 
 
-@pytest.mark.slow
+@pytest.mark.pymc
+@pytest.mark.flow
 def test_normalizing_flow():
     with pm.Model() as model:
         pm.HalfNormal("x", shape=2)
@@ -272,6 +297,7 @@ def test_normalizing_flow():
     assert kstest.pvalue > 0.01
 
 
+@pytest.mark.pymc
 @pytest.mark.parametrize(
     ("backend", "gradient_backend"),
     [

tests/test_pymc.py

@@ -1,9 +1,16 @@
+from importlib.util import find_spec
+import pytest
+
+if find_spec("bridgestan") is None:
+    pytest.skip("Skip stan tests", allow_module_level=True)
+
 import numpy as np
 import pytest
 
 import nutpie
 
 
+@pytest.mark.stan
 def test_stan_model():
     model = """
     data {}
@@ -20,6 +27,7 @@ def test_stan_model():
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.stan
 def test_stan_model_data():
     model = """
     data {
@@ -40,11 +48,55 @@ def test_stan_model_data():
     trace.posterior.a  # noqa: B018
 
 
-@pytest.mark.slow
-def test_stan_flow():
+@pytest.mark.stan
+def test_stan_memory_order():
     model = """
+    data {
+        real x;
+    }
     parameters {
         real a;
+    }
+    model {
+        a ~ normal(0, 1);
+    }
+    generated quantities {
+        array[2, 3] matrix[5, 7] b;
+        real count = 0;
+        for (i in 1:2)
+            for (j in 1:3) {
+                for (k in 1:5) {
+                    for (n in 1:7) {
+                        b[i, j][k, n] = count;
+                        count = count + 1;
+                    }
+                }
+            }
+        }
+    """
+
+    compiled_model = nutpie.compile_stan_model(code=model)
+    with pytest.raises(RuntimeError):
+        trace = nutpie.sample(compiled_model)
+    trace = nutpie.sample(compiled_model.with_data(x=np.array(3.0)))
+    trace.posterior.a  # noqa: B018
+    assert trace.posterior.b.shape == (6, 1000, 2, 3, 5, 7)
+    b = trace.posterior.b.isel(chain=0, draw=0)
+    count = 0
+    for i in range(2):
+        for j in range(3):
+            for k in range(5):
+                for n in range(7):
+                    assert float(b[i, j, k, n]) == count
+                    count += 1
+
+
+@pytest.mark.flow
+@pytest.mark.stan
+def test_stan_flow():
+    model = """
+    parameters {
+        array[5] real a;
         real<lower=0> b;
     }
     model {