Better error message for NaN occurred in optimization

RELEASE-NOTES.md

@@ -9,7 +9,8 @@
 - Add `check_test_point` method to `pm.Model`
 - Add `Ordered` Transformation and `OrderedLogistic` distribution
 - Add `Chain` transformation
-- Better warning message for `Mass matrix contains zeros on the diagonal. Some derivatives might always be zero`
+- Improve error message `Mass matrix contains zeros on the diagonal. Some derivatives might always be zero` during tuning of `pm.sample`
+- Improve error message `NaN occurred in optimization.` during ADVI
 - Save and load traces without `pickle` using `pm.save_trace` and `pm.load_trace`
 
 ### Fixes

pymc3/variational/inference.py

@@ -147,8 +147,24 @@ def _iterate_without_loss(self, s, _, step_func, progress, callbacks):
         try:
             for i in progress:
                 step_func()
-                if np.isnan(self.approx.params[0].get_value()).any():
-                    raise FloatingPointError('NaN occurred in optimization.')
+                current_param = self.approx.params[0].get_value()
+                if np.isnan(current_param).any():
+                    name_slc = []
+                    tmp_hold = list(range(current_param.size))
+                    vmap = self.approx.groups[0].bij.ordering.vmap
+                    for vmap_ in vmap:
+                        slclen = len(tmp_hold[vmap_.slc])
+                        for i in range(slclen):
+                            name_slc.append((vmap_.var, i))
+                    index = np.where(np.isnan(current_param))[0]
+                    errmsg = ['NaN occurred in optimization. ']
+                    suggest_solution = 'Try tracking this parameter: ' \
+                                       'http://docs.pymc.io/notebooks/variational_api_quickstart.html#Tracking-parameters'
+                    for ii in index:
+                        errmsg.append('The current approximation of RV `{}`.ravel()[{}]'
+                                      ' is NaN.'.format(*name_slc[ii]))
+                    errmsg.append(suggest_solution)
+                    raise FloatingPointError('\n'.join(errmsg))
                 for callback in callbacks:
                     callback(self.approx, None, i+s+1)
         except (KeyboardInterrupt, StopIteration) as e:
@@ -178,7 +194,23 @@ def _infmean(input_array):
                 if np.isnan(e):  # pragma: no cover
                     scores = scores[:i]
                     self.hist = np.concatenate([self.hist, scores])
-                    raise FloatingPointError('NaN occurred in optimization.')
+                    current_param = self.approx.params[0].get_value()
+                    name_slc = []
+                    tmp_hold = list(range(current_param.size))
+                    vmap = self.approx.groups[0].bij.ordering.vmap
+                    for vmap_ in vmap:
+                        slclen = len(tmp_hold[vmap_.slc])
+                        for i in range(slclen):
+                            name_slc.append((vmap_.var, i))
+                    index = np.where(np.isnan(current_param))[0]
+                    errmsg = ['NaN occurred in optimization. ']
+                    suggest_solution = 'Try tracking this parameter: ' \
+                                       'http://docs.pymc.io/notebooks/variational_api_quickstart.html#Tracking-parameters'
+                    for ii in index:
+                        errmsg.append('The current approximation of RV `{}`.ravel()[{}]'
+                                      ' is NaN.'.format(*name_slc[ii]))
+                    errmsg.append(suggest_solution)
+                    raise FloatingPointError('\n'.join(errmsg))
                 scores[i] = e
                 if i % 10 == 0:
                     avg_loss = _infmean(scores[max(0, i - 1000):i + 1])