Fix seeding bug in sample_numpyro_nuts when more than one chain was sampled

Author: ricardoV94

pymc/sampling_jax.py

@@ -152,10 +152,8 @@ def _get_batched_jittered_initial_points(
         list with one item per variable and number of chains as batch dimension.
         Each item has shape `(chains, *var.shape)`
     """
-    if isinstance(random_seed, (int, np.integer)):
-        random_seed = np.random.default_rng(random_seed).integers(2**30, size=chains)
-    elif not isinstance(random_seed, (list, tuple, np.ndarray)):
-        raise ValueError(f"The `seeds` must be int or array-like. Got {type(random_seed)} instead.")
+
+    random_seed = np.random.default_rng(random_seed).integers(2**30, size=chains)
 
     assert len(random_seed) == chains
 
@@ -213,9 +211,7 @@ def sample_numpyro_nuts(
         dims = {}
 
     if random_seed is None:
-        random_seed = model.rng_seeder.randint(
-            2**30, dtype=np.int64, size=chains if chains > 1 else None
-        )
+        random_seed = model.rng_seeder.randint(2**30, dtype=np.int64)
 
     tic1 = datetime.now()
     print("Compiling...", file=sys.stdout)

pymc/tests/test_sampling_jax.py

@@ -161,3 +161,32 @@ def test_get_batched_jittered_initial_points():
 
     assert ips[0].shape == (2, 2, 3)
     assert np.all(ips[0][0] != ips[0][1])
+
+
+@pytest.mark.parametrize("random_seed", (None, 123))
+@pytest.mark.parametrize("chains", (1, 2))
+def test_seeding(chains, random_seed):
+    sample_kwargs = dict(
+        tune=100,
+        draws=5,
+        chains=chains,
+        random_seed=random_seed,
+    )
+
+    with pm.Model(rng_seeder=456) as m:
+        pm.Normal("x", mu=0, sigma=1)
+        result1 = sample_numpyro_nuts(**sample_kwargs)
+
+    with pm.Model(rng_seeder=456) as m:
+        pm.Normal("x", mu=0, sigma=1)
+        result2 = sample_numpyro_nuts(**sample_kwargs)
+        result3 = sample_numpyro_nuts(**sample_kwargs)
+
+    assert np.all(result1.posterior["x"] == result2.posterior["x"])
+    expected_equal_result3 = random_seed is not None
+    assert np.all(result2.posterior["x"] == result3.posterior["x"]) == expected_equal_result3
+
+    if chains > 1:
+        assert np.all(result1.posterior["x"].sel(chain=0) != result1.posterior["x"].sel(chain=1))
+        assert np.all(result2.posterior["x"].sel(chain=0) != result2.posterior["x"].sel(chain=1))
+        assert np.all(result3.posterior["x"].sel(chain=0) != result3.posterior["x"].sel(chain=1))