test: check that normalizing flows are reproducible

Author: aseyboldt

tests/test_pymc.py

@@ -270,7 +270,7 @@ def test_pymc_var_names(backend, gradient_backend):
 
 @pytest.mark.pymc
 @pytest.mark.flow
-@pytest.mark.parametrize("kind", ["masked", "subset"])
+@pytest.mark.parametrize("kind", ["masked"])
 def test_normalizing_flow(kind):
     with pm.Model() as model:
         pm.HalfNormal("x", shape=2)
@@ -280,7 +280,7 @@ def test_normalizing_flow(kind):
     ).with_transform_adapt(
         verbose=True,
         coupling_type=kind,
-        num_layers=4,
+        num_layers=2,
     )
     trace = nutpie.sample(
         compiled,
@@ -290,13 +290,38 @@ def test_normalizing_flow(kind):
         seed=1,
         draws=2000,
     )
-    draws = trace.posterior.x.isel(x_dim_0=0, chain=0)
-    kstest = stats.ks_1samp(draws, stats.halfnorm.cdf)
-    assert kstest.pvalue > 0.01
 
-    draws = trace.posterior.x.isel(x_dim_0=1, chain=0)
-    kstest = stats.ks_1samp(draws, stats.halfnorm.cdf)
-    assert kstest.pvalue > 0.01
+    compiled = nutpie.compile_pymc_model(
+        model, backend="jax", gradient_backend="jax"
+    ).with_transform_adapt(
+        verbose=True,
+        coupling_type=kind,
+        num_layers=2,
+    )
+    trace2 = nutpie.sample(
+        compiled,
+        chains=1,
+        transform_adapt=True,
+        window_switch_freq=128,
+        seed=1,
+        draws=2000,
+    )
+    draws1 = trace.posterior.x
+    draws2 = trace2.posterior.x
+
+    # Check that the two draws are the same
+    assert np.allclose(draws1, draws2)
+
+    # Compare to precompute values to make sure it is reproducible
+    # accross architectures
+    expected = np.array([
+        [1.81033486, 1.18735544],
+        [0.12551686, 0.04161655],
+        [1.07813544, 0.12578679],
+        [0.71503155, 0.37380833],
+        [0.83237662, 0.67041153]
+    ])
+    assert np.allclose(draws1.isel(chain=0, draw=slice(0, 5)), expected, atol=1e-5)
 
 
 @pytest.mark.pymc