fix test after latest releases jax 0.2.9 and jaxlib 0.1.59

Author: mganahl

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):