Repository: google/jax

Issue:
jax.numpy.linalg.eig - 'jax_debug_nans' related bug
Reproducing the error with a simple matrix: (same issue with `np.linalg.eigh`)
```
A = np.array([[1.,2.],[2.,3.]])
D, V = np.linalg.eig(A)
```
Output
```
TypeError                                 Traceback (most recent call last)
<ipython-input-171-c4c8dbbffe17> in <module>
      1 A = np.array([[1.,2.],[2.,3.]])
----> 2 V,D = np.linalg.eig(A)

~/RESEARCH/jax/jax/numpy/linalg.py in eig(a)
     95 def eig(a):
     96   a = _promote_arg_dtypes(np.asarray(a))
---> 97   w, vl, vr = lax_linalg.eig(a)
     98   return w, vr
     99 

~/RESEARCH/jax/jax/lax_linalg.py in eig(x)
     47 
     48 def eig(x):
---> 49   w, vl, vr = eig_p.bind(x)
     50   return w, vl, vr
     51 

~/RESEARCH/jax/jax/core.py in bind(self, *args, **kwargs)
    128     top_trace = find_top_trace(args)
    129     if top_trace is None:
--> 130       return self.impl(*args, **kwargs)
    131 
    132     tracers = map(top_trace.full_raise, args)

~/RESEARCH/jax/jax/lax_linalg.py in eig_impl(operand)
    153 
    154 def eig_impl(operand):
--> 155   return xla.apply_primitive(eig_p, operand)
    156 
    157 def eig_translation_rule(c, operand):

~/RESEARCH/jax/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    122   abstract_args = map(abstractify, args)
    123   compiled_fun = xla_primitive_callable(prim, *abstract_args, **params)
--> 124   return compiled_fun(*args)
    125 
    126 @cache()

~/RESEARCH/jax/jax/interpreters/xla.py in _execute_compiled_primitive(prim, compiled, backend, result_handler, *args)
    172   input_bufs = [device_put(x, device_num, backend=backend) for x in args]
    173   out_buf = compiled.Execute(input_bufs)
--> 174   if FLAGS.jax_debug_nans: check_nans(prim, out_buf)
    175   return result_handler(out_buf)
    176 

~/RESEARCH/jax/jax/interpreters/xla.py in check_nans(prim, bufs)
    177 def check_nans(prim, bufs):
    178   if prim.multiple_results:
--> 179     for buf in bufs:
    180       _check_nans(prim.name, buf.shape(), buf)
    181   else:

TypeError: 'jaxlib.xla_extension.PyLocalBuffer' object is not iterable
```


Discussion:
UPDATE  (bug still persists)
---

The bug is related to the flag `jax_debug_nans` that I have set `True` as default in my workflow. If set it otherwise, an error is not thrown. However, I need this flag set for debugging reasons. Is it a quick fix?

```
from jax.config import config
config.update("jax_debug_nans", False)
A = np.array([[1.,2.],[2.,3.]])
D, V = np.linalg.eig(A)
D, V
```
Output
```
(DeviceArray([-0.23606798+0.j,  4.23606798+0.j]),
 DeviceArray([[-0.85065081+0.j, -0.52573111+0.j],
              [ 0.52573111+0.j, -0.85065081+0.j]]))
```


Resolving patch:
diff --git a/jax/config.py b/jax/config.py
--- a/jax/config.py
+++ b/jax/config.py
@@ -23,10 +23,13 @@ def __init__(self):
     self.use_absl = False
 
   def update(self, name, val):
-    self.check_exists(name)
-    if name not in self.values:
-      raise Exception("Unrecognized config option: {}".format(name))
-    self.values[name] = val
+    if self.use_absl:
+      setattr(self.absl_flags.FLAGS, name, val)
+    else:
+      self.check_exists(name)
+      if name not in self.values:
+        raise Exception("Unrecognized config option: {}".format(name))
+      self.values[name] = val
 
   def read(self, name):
     if self.use_absl:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -171,7 +171,8 @@ def _execute_compiled_primitive(prim, compiled, backend, result_handler, *args):
   device_num, = compiled.DeviceOrdinals()
   input_bufs = [device_put(x, device_num, backend=backend) for x in args]
   out_buf = compiled.Execute(input_bufs)
-  if FLAGS.jax_debug_nans: check_nans(prim, out_buf)
+  if FLAGS.jax_debug_nans:
+    check_nans(prim, out_buf.destructure() if prim.multiple_results else out_buf)
   return result_handler(out_buf)
 
 def check_nans(prim, bufs):


Test patch:
diff --git a/tests/debug_nans_test.py b/tests/debug_nans_test.py
new file mode 100644
--- /dev/null
+++ b/tests/debug_nans_test.py
@@ -0,0 +1,58 @@
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for --debug_nans."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+import numpy as onp
+
+import jax
+from jax import test_util as jtu
+from jax.test_util import check_grads
+from jax import numpy as np
+from jax import random
+
+from jax.config import config
+config.parse_flags_with_absl()
+
+class DebugNaNsTest(jtu.JaxTestCase):
+
+  def setUp(self):
+    self.cfg = config.read("jax_debug_nans")
+    config.update("jax_debug_nans", True)
+
+  def tearDown(self):
+    config.update("jax_debug_nans", self.cfg)
+
+  def testSingleResultPrimitiveNoNaN(self):
+    A = np.array([[1., 2.], [2., 3.]])
+    B = np.tanh(A)
+
+  def testMultipleResultPrimitiveNoNaN(self):
+    A = np.array([[1., 2.], [2., 3.]])
+    D, V = np.linalg.eig(A)
+
+  def testJitComputationNoNaN(self):
+    A = np.array([[1., 2.], [2., 3.]])
+    B = jax.jit(np.tanh)(A)
+
+  def testSingleResultPrimitiveNaN(self):
+    A = np.array(0.)
+    with self.assertRaises(FloatingPointError):
+      B = 0. / A