fix(stan): rng for generated quantities

Generated quantities have their own random number generator.
The seed for this generator did not depend on the global
seed of the model, so that random stream for two different
sampler runs were reusing the same randomness.

Author: aseyboldt

src/stan.rs

@@ -531,8 +531,8 @@ impl Model for StanModel {
 
     fn new_trace<'a, S: Settings, R: rand::Rng + ?Sized>(
         &'a self,
-        _rng: &mut R,
-        chain: u64,
+        rng: &mut R,
+        _chain: u64,
         settings: &S,
     ) -> anyhow::Result<Self::DrawStorage<'a, S>> {
         let draws = settings.hint_num_tune() + settings.hint_num_draws();
@@ -541,7 +541,8 @@ impl Model for StanModel {
             .iter()
             .map(|var| Vec::with_capacity(var.size * draws))
             .collect();
-        let rng = self.model.new_rng(chain as u32)?;
+        let seed = rng.next_u32();
+        let rng = self.model.new_rng(seed)?;
         let buffer = vec![0f64; self.model.param_num(true, true)];
         Ok(StanTrace {
             model: self,

tests/test_stan.py

@@ -27,6 +27,43 @@ def test_stan_model():
     trace.posterior.a  # noqa: B018
 
 
+@pytest.mark.stan
+def test_seed():
+    model = """
+    data {}
+    parameters {
+        real a;
+    }
+    model {
+        a ~ normal(0, 1);
+    }
+    generated quantities {
+        real b = normal_rng(0, 1);
+    }
+    """
+
+    compiled_model = nutpie.compile_stan_model(code=model)
+    trace = nutpie.sample(compiled_model, seed=42)
+    trace2 = nutpie.sample(compiled_model, seed=42)
+    trace3 = nutpie.sample(compiled_model, seed=43)
+
+    assert np.allclose(trace.posterior.a, trace2.posterior.a)
+    assert np.allclose(trace.posterior.b, trace2.posterior.b)
+
+    assert not np.allclose(trace.posterior.a, trace3.posterior.a)
+    assert not np.allclose(trace.posterior.b, trace3.posterior.b)
+    # Check that all chains are pairwise different
+    for i in range(len(trace.posterior.a)):
+        for j in range(i + 1, len(trace.posterior.a)):
+            assert not np.allclose(trace.posterior.a[i], trace.posterior.a[j])
+            assert not np.allclose(trace.posterior.b[i], trace.posterior.b[j])
+    # Check that all chains are pairwise different between seeds
+    for i in range(len(trace.posterior.a)):
+        for j in range(len(trace3.posterior.a)):
+            assert not np.allclose(trace.posterior.a[i], trace3.posterior.a[j])
+            assert not np.allclose(trace.posterior.b[i], trace3.posterior.b[j])
+
+
 @pytest.mark.stan
 def test_stan_model_data():
     model = """