test: check that normalizing flows are reproducible

Author: aseyboldt

python/nutpie/normalizing_flow.py

@@ -1916,7 +1916,7 @@ def make_flow(
         diag = jnp.sqrt(pos_std / grad_std)
         mean = positions.mean(0) + gradients.mean(0) * diag * diag
 
-    key = jax.random.PRNGKey(seed % (2**63))
+    key = jax.random.key(seed % (2**63), impl="threefry2x32")
 
     diag_param = Parameterize(
         lambda x: x + jnp.sqrt(1 + x**2),

tests/test_pymc.py

@@ -8,7 +8,6 @@
 import numpy as np
 import pymc as pm
 import pytest
-from scipy import stats
 
 import nutpie
 import nutpie.compile_pymc
@@ -270,38 +269,76 @@ 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)
+    import jax
 
-    compiled = nutpie.compile_pymc_model(
-        model, backend="jax", gradient_backend="jax"
-    ).with_transform_adapt(
-        verbose=True,
-        coupling_type=kind,
-        num_layers=4,
-    )
-    trace = nutpie.sample(
-        compiled,
-        chains=1,
-        transform_adapt=True,
-        window_switch_freq=128,
-        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
+    old_x64 = jax.config.update("jax_enable_x64", True)
+
+    try:
+        with pm.Model() as model:
+            pm.HalfNormal("x", shape=2)
 
-    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,
+        )
+        trace = nutpie.sample(
+            compiled,
+            chains=1,
+            transform_adapt=True,
+            window_switch_freq=128,
+            seed=1,
+            draws=2000,
+        )
+
+        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
+        np.testing.assert_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],
+            ]
+        )
+        print(expected)
+        np.testing.assert_allclose(
+            draws1.isel(chain=0, draw=slice(0, 5)).values, expected, atol=1e-5
+        )
+    finally:
+        # Restore the original config
+        jax.config.update("jax_enable_x64", old_x64)
 
 
 @pytest.mark.pymc
 @pytest.mark.flow
-@pytest.mark.parametrize("kind", ["masked", "subset"])
+@pytest.mark.parametrize("kind", ["masked"])
 def test_normalizing_flow_1d(kind):
     with pm.Model() as model:
         pm.HalfNormal("x")