test: add low rank tests

Author: aseyboldt

tests/test_pymc.py

@@ -31,6 +31,32 @@ def test_pymc_model(backend, gradient_backend):
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.pymc
+@parameterize_backends
+def test_low_rank(backend, gradient_backend):
+    with pm.Model() as model:
+        pm.Normal("a")
+
+    compiled = nutpie.compile_pymc_model(
+        model, backend=backend, gradient_backend=gradient_backend
+    )
+    trace = nutpie.sample(compiled, chains=1, low_rank_modified_mass_matrix=True)
+    trace.posterior.a  # noqa: B018
+
+
+@pytest.mark.pymc
+@parameterize_backends
+def test_low_rank_half_normal(backend, gradient_backend):
+    with pm.Model() as model:
+        pm.HalfNormal("a", shape=13)
+
+    compiled = nutpie.compile_pymc_model(
+        model, backend=backend, gradient_backend=gradient_backend
+    )
+    trace = nutpie.sample(compiled, chains=1, low_rank_modified_mass_matrix=True)
+    trace.posterior.a  # noqa: B018
+
+
 @pytest.mark.pymc
 @parameterize_backends
 def test_zero_size(backend, gradient_backend):

tests/test_stan.py

@@ -27,6 +27,23 @@ def test_stan_model():
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.stan
+def test_stan_model_low_rank():
+    model = """
+    data {}
+    parameters {
+        real a;
+    }
+    model {
+        a ~ normal(0, 1);
+    }
+    """
+
+    compiled_model = nutpie.compile_stan_model(code=model)
+    trace = nutpie.sample(compiled_model, low_rank_modified_mass_matrix=True)
+    trace.posterior.a  # noqa: B018
+
+
 @pytest.mark.stan
 def test_empty():
     model = """
@@ -40,7 +57,7 @@ def test_empty():
     """
 
     compiled_model = nutpie.compile_stan_model(code=model)
-    trace = nutpie.sample(compiled_model)  # noqa: F841
+    nutpie.sample(compiled_model)
     # TODO: Variable `a` is missing because of this bridgestan issue:
     # https://github.com/roualdes/bridgestan/issues/278
     # assert trace.posterior.a.shape == (0, 1000)