Merge branch 'master' into master

google · Jan 27, 2021 · a8b294f · a8b294f
2 parents ba88d50 + 72bb96c
commit a8b294f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/requirements_travis.txt b/requirements_travis.txt
@@ -2,5 +2,5 @@ tensorflow>=2.0.0
 pytest
 torch>=1.4.0
 jax>=0.1.68
-jaxlib==0.1.57
+jaxlib>=0.1.59
 pylint==2.5.3
diff --git a/tensornetwork/tests/ncon_interface_test.py b/tensornetwork/tests/ncon_interface_test.py
@@ -56,9 +56,12 @@ def test_return_type(backend):
   result_1 = ncon_interface.ncon([t1, t2], [(-1, 1), (1, -2)], backend=backend)
   result_2 = ncon_interface.ncon([n1, n2], [(-1, 1), (1, -2)], backend=backend)
   result_3 = ncon_interface.ncon([n1, t2], [(-1, 1), (1, -2)], backend=backend)
-  assert isinstance(result_1, type(n1.backend.convert_to_tensor(t1)))
   assert isinstance(result_2, Tensor)
-  assert isinstance(result_3, type(n1.backend.convert_to_tensor(t1)))
+  if backend not in ('jax', get_backend('jax')):
+    # jitted functions return jaxlib.xla_extension.Buffer,
+    # convert_to_tensor returns jax.interpreters.xla._DeviceArray now.
+    assert isinstance(result_1, type(n1.backend.convert_to_tensor(t1)))
+    assert isinstance(result_3, type(n1.backend.convert_to_tensor(t1)))
 
 
 def test_order_spec(backend):