fix test catch multable state warning

pyro-ppl · Oct 20, 2023 · f577541 · f577541
1 parent 5ccda17
commit f577541
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/numpyro/infer/elbo.py b/numpyro/infer/elbo.py
@@ -214,7 +214,7 @@ def get_model_density(key, latent):
                     return elbo_particle, mutable_params
                 else:
                     warnings.warn(
-                        "mutable parameters is currently ignored when num_particles > 1."
+                        "mutable state is currently ignored when num_particles > 1."
                     )
                     return elbo_particle, None
             else:
@@ -347,9 +347,10 @@ def single_particle_elbo(rng_key):
                 if self.num_particles == 1:
                     return elbo_particle, mutable_params
                 else:
-                    raise ValueError(
-                        "Currently, we only support mutable states with num_particles=1."
+                    warnings.warn(
+                        "mutable state is currently ignored when num_particles > 1."
                     )
+                    return elbo_particle, None
             else:
                 return elbo_particle, None
 

diff --git a/test/infer/test_svi.py b/test/infer/test_svi.py
@@ -463,7 +463,7 @@ def guide():
 
     svi = SVI(model, guide, optim.Adam(0.1), elbo(num_particles=num_particles))
     if num_particles > 1:
-        with pytest.raises(ValueError, match="mutable state"):
+        with pytest.warns(UserWarning, match="mutable state"):
             svi_result = svi.run(random.PRNGKey(0), 1000, stable_update=stable_update)
         return
     svi_result = svi.run(random.PRNGKey(0), 1000, stable_update=stable_update)