Fix ensemble mcmc multiple run

numpyro/infer/ensemble.py

@@ -154,15 +154,8 @@ def init(
             "EnsembleSampler only supports chain_method='vectorized' with num_chains > 1.\n"
             "If you want to run chains in parallel, please raise a github issue."
         )
-
         assert rng_key.shape[0] % 2 == 0, "Number of chains must be even."
-
         self._num_chains = rng_key.shape[0]
-
-        if self._potential_fn and init_params is None:
-            raise ValueError(
-                "Valid value of `init_params` must be provided with `potential_fn`."
-            )
         if init_params is not None:
             assert all(
                 [
@@ -178,6 +171,11 @@ def init(
             rng_key_init_model, model_args, model_kwargs, init_params
         )
 
+        if self._potential_fn and init_params is None:
+            raise ValueError(
+                "Valid value of `init_params` must be provided with `potential_fn`."
+            )
+
         self._num_warmup = num_warmup
 
         return EnsembleSamplerState(

test/infer/test_ensemble_mcmc.py

@@ -79,3 +79,20 @@ def test_out_shape_smoke(kernel_cls):
 def test_invalid_moves(kernel_cls):
     with pytest.raises(AssertionError, match="Each move"):
         kernel_cls(model, moves={"invalid": 1.0})
+
+
+@pytest.mark.parametrize("kernel_cls", [AIES, ESS])
+def test_multirun(kernel_cls):
+    n_chains = 10
+    kernel = kernel_cls(model)
+
+    mcmc = MCMC(
+        kernel,
+        num_warmup=10,
+        num_samples=10,
+        progress_bar=False,
+        num_chains=n_chains,
+        chain_method="vectorized",
+    )
+    mcmc.run(random.PRNGKey(2), labels)
+    mcmc.run(random.PRNGKey(3), labels)