Bug/steinvi reinit (#1626)

* added separate guide for reinitialization.

* added test case for reinit.

numpyro/contrib/einstein/steinvi.py

@@ -137,6 +137,7 @@ def __init__(
         self._inference_model = model
         self.model = model
         self.guide = guide
+        self._init_guide = deepcopy(guide)
         self.optim = optim
         self.stein_loss = SteinLoss(  # TODO: @OlaRonning handle enum
             elbo_num_particles=num_elbo_particles,
@@ -388,7 +389,7 @@ def init(self, rng_key: KeyArray, *args, **kwargs):
         )
 
         guide_init_params = self._find_init_params(
-            particle_seed, self.guide, args, kwargs
+            particle_seed, self._init_guide, args, kwargs
         )
 
         guide_init = handlers.seed(self.guide, guide_seed)

test/contrib/einstein/test_steinvi.py

@@ -10,7 +10,7 @@
 from numpy.testing import assert_allclose
 import pytest
 
-from jax import random
+from jax import numpy as jnp, random
 
 import numpyro
 from numpyro import handlers
@@ -193,6 +193,27 @@ def model():
         return
 
 
+def test_stein_reinit():
+    num_particles = 4
+
+    def model():
+        numpyro.sample("x", Normal(0, 1.0))
+
+    stein = SteinVI(
+        model,
+        AutoDelta(model),
+        Adam(1.0),
+        RBFKernel(),
+        num_stein_particles=num_particles,
+    )
+
+    for i in range(2):
+        with handlers.seed(rng_seed=i):
+            params = stein.get_params(stein.init(numpyro.prng_key()))
+            xs = params["x_auto_loc"]
+            assert jnp.unique(xs).shape == xs.shape
+
+
 @pytest.mark.parametrize(
     "auto_class",
     [