Repository: pyro-ppl/numpyro

Issue:
Skip current update in SVI if loss gets nan
cc @martinjankowiak 
Unhelpful Error Dirichlet Process Mixture Model
Following my post here, https://forum.pyro.ai/t/dirichlet-process-gaussian-mixture-model-typeerror-for-small-concentration-parameter/2685, when one tries creating a Dirichlet Process Mixture Model with 0 mixture components, the error message returned is 

`TypeError: iota shape must have every element be nonnegative, got (-1,).`

This isn't helpful in identifying the source of the error. Could a better message be returned?


Discussion:
hmm actually i think it'd be more useful to skip the update if any of the gradients are nans (you can have nan gradients but a well-defined loss)
yes, we will skip if there is any nan/inf on params values
Thank you, @RylanSchaeffer. I think all we need is to add an assertation `size > 0` in the [constructor](https://github.com/pyro-ppl/numpyro/blob/e2e82189e96f45a1d2715c15217b85ed55e9f16e/numpyro/primitives.py#L347) of `plate`.

Resolving patch:
diff --git a/numpyro/infer/svi.py b/numpyro/infer/svi.py
--- a/numpyro/infer/svi.py
+++ b/numpyro/infer/svi.py
@@ -146,7 +146,25 @@ def update(self, svi_state, *args, **kwargs):
         loss_val, optim_state = self.optim.eval_and_update(loss_fn, svi_state.optim_state)
         return SVIState(optim_state, rng_key), loss_val
 
-    def run(self, rng_key, num_steps, *args, progress_bar=True, **kwargs):
+    def stable_update(self, svi_state, *args, **kwargs):
+        """
+        Similar to :meth:`update` but returns the current state if the
+        the loss or the new state contains invalid values.
+
+        :param svi_state: current state of SVI.
+        :param args: arguments to the model / guide (these can possibly vary during
+            the course of fitting).
+        :param kwargs: keyword arguments to the model / guide (these can possibly vary
+            during the course of fitting).
+        :return: tuple of `(svi_state, loss)`.
+        """
+        rng_key, rng_key_step = random.split(svi_state.rng_key)
+        loss_fn = partial(_apply_loss_fn, self.loss.loss, rng_key_step, self.constrain_fn, self.model,
+                          self.guide, args, kwargs, self.static_kwargs)
+        loss_val, optim_state = self.optim.eval_and_stable_update(loss_fn, svi_state.optim_state)
+        return SVIState(optim_state, rng_key), loss_val
+
+    def run(self, rng_key, num_steps, *args, progress_bar=True, stable_update=False, **kwargs):
         """
         (EXPERIMENTAL INTERFACE) Run SVI with `num_steps` iterations, then return
         the optimized parameters and the stacked losses at every step. If `num_steps`
@@ -162,6 +180,8 @@ def run(self, rng_key, num_steps, *args, progress_bar=True, **kwargs):
         :param args: arguments to the model / guide
         :param bool progress_bar: Whether to enable progress bar updates. Defaults to
             ``True``.
+        :param bool stable_update: whether to use :meth:`stable_update` to update
+            the state. Defaults to False.
         :param kwargs: keyword arguments to the model / guide
         :return: a namedtuple with fields `params` and `losses` where `params`
             holds the optimized values at :class:`numpyro.param` sites,
@@ -169,7 +189,10 @@ def run(self, rng_key, num_steps, *args, progress_bar=True, **kwargs):
         :rtype: SVIRunResult
         """
         def body_fn(svi_state, carry):
-            svi_state, loss = self.update(svi_state, *args, **kwargs)
+            if stable_update:
+                svi_state, loss = self.stable_update(svi_state, *args, **kwargs)
+            else:
+                svi_state, loss = self.update(svi_state, *args, **kwargs)
             return svi_state, loss
 
         svi_state = self.init(rng_key, *args, **kwargs)
@@ -181,7 +204,15 @@ def body_fn(svi_state, carry):
                     svi_state, loss = jit(body_fn)(svi_state, None)
                     losses.append(loss)
                     if i % batch == 0:
-                        avg_loss = sum(losses[i-batch:]) / batch
+                        if stable_update:
+                            valid_losses = [x for x in losses[i - batch:] if x == x]
+                            num_valid = len(valid_losses)
+                            if num_valid == 0:
+                                avg_loss = float('nan')
+                            else:
+                                avg_loss = sum(valid_losses) / num_valid
+                        else:
+                            avg_loss = sum(losses[i-batch:]) / batch
                         t.set_postfix_str("init loss: {:.4f}, avg. loss [{}-{}]: {:.4f}"
                                           .format(losses[0], i - batch + 1, i, avg_loss),
                                           refresh=False)
diff --git a/numpyro/optim.py b/numpyro/optim.py
--- a/numpyro/optim.py
+++ b/numpyro/optim.py
@@ -10,7 +10,7 @@
 from collections import namedtuple
 from typing import Callable, Tuple, TypeVar
 
-from jax import value_and_grad
+from jax import lax, value_and_grad
 from jax.experimental import optimizers
 from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
@@ -77,6 +77,24 @@ def eval_and_update(self, fn: Callable, state: _IterOptState) -> _IterOptState:
         out, grads = value_and_grad(fn)(params)
         return out, self.update(grads, state)
 
+    def eval_and_stable_update(self, fn: Callable, state: _IterOptState) -> _IterOptState:
+        """
+        Like :meth:`eval_and_update` but when the value of the objective function
+        or the gradients are not finite, we will not update the input `state`
+        and will set the objective output to `nan`.
+
+        :param fn: objective function.
+        :param state: current optimizer state.
+        :return: a pair of the output of objective function and the new optimizer state.
+        """
+        params = self.get_params(state)
+        out, grads = value_and_grad(fn)(params)
+        out, state = lax.cond(jnp.isfinite(out) & jnp.isfinite(ravel_pytree(grads)[0]).all(),
+                              lambda _: (out, self.update(grads, state)),
+                              lambda _: (jnp.nan, state),
+                              None)
+        return out, state
+
     def get_params(self, state: _IterOptState) -> _Params:
         """
         Get current parameter values.
diff --git a/numpyro/primitives.py b/numpyro/primitives.py
--- a/numpyro/primitives.py
+++ b/numpyro/primitives.py
@@ -344,6 +344,7 @@ class plate(Messenger):
 
     def __init__(self, name, size, subsample_size=None, dim=None):
         self.name = name
+        assert size > 0, "size of plate should be positive"
         self.size = size
         if dim is not None and dim >= 0:
             raise ValueError('dim arg must be negative.')


Test patch:
diff --git a/test/infer/test_reparam.py b/test/infer/test_reparam.py
--- a/test/infer/test_reparam.py
+++ b/test/infer/test_reparam.py
@@ -124,7 +124,7 @@ def test_reparam_log_joint(model, kwargs):
     latent_y = neutra.transform(latent_x)
     log_det_jacobian = neutra.transform.log_abs_det_jacobian(latent_x, latent_y)
     pe = pe_fn(guide._unpack_latent(latent_y))
-    assert_allclose(pe_transformed, pe - log_det_jacobian)
+    assert_allclose(pe_transformed, pe - log_det_jacobian, rtol=2e-7)
 
 
 @pytest.mark.parametrize("shape", [(), (4,), (3, 2)], ids=str)
diff --git a/test/infer/test_svi.py b/test/infer/test_svi.py
--- a/test/infer/test_svi.py
+++ b/test/infer/test_svi.py
@@ -186,3 +186,19 @@ def guide():
 
     svi = SVI(model, guide, optim.Adam(1), Trace_ELBO())
     svi.run(random.PRNGKey(0), num_steps)
+
+
+@pytest.mark.parametrize("stable_run", [True, False])
+def test_stable_run(stable_run):
+
+    def model():
+        var = numpyro.sample("var", dist.Exponential(1))
+        numpyro.sample("obs", dist.Normal(0, jnp.sqrt(var)), obs=0.)
+
+    def guide():
+        loc = numpyro.param("loc", 0.)
+        numpyro.sample("var", dist.Normal(loc, 10))
+
+    svi = SVI(model, guide, optim.Adam(1), Trace_ELBO())
+    svi_result = svi.run(random.PRNGKey(0), 1000, stable_update=stable_run)
+    assert jnp.isfinite(svi_result.params["loc"]) == stable_run