Merge pull request #1090 from hawkinsp/pytree

Fix test failures.

jax/api.py

@@ -1109,8 +1109,10 @@ def _device_get(x):
 
 def device_get(x):
   for y in tree_leaves(x):
-    if not isinstance(y, core.Tracer):
+    try:
       y.copy_to_host_async()
+    except AttributeError:
+      pass
   return tree_map(_device_get, x)
 
 

tests/core_test.py

@@ -184,7 +184,7 @@ class CoreTest(jtu.JaxTestCase):
     try:
       tree_multimap(f, xs, ys_bad)
       assert False
-    except TypeError:
+    except (TypeError, ValueError):
       pass
 
   def test_print_jaxpr_compound(self):