Temporarily disable test_jit_device_assignment.

Future breaking changes to jaxlib will break this test, so disable it
until we update jaxlib and then can update jax and reenable the test.

tests/api_test.py

@@ -920,6 +920,7 @@ def scalar_solve2(f, y):
     self.assertAllClose(grad, api.grad(pow)(5.0, 1.5), check_dtypes=False)
 
   def test_jit_device_assignment(self):
+    raise unittest.SkipTest("Temporarily disabled while device API is being changed.")
     device_num = xb.device_count() - 1
     x = api.jit(lambda x: x, device_assignment=device_num)(3.)
     self.assertIsInstance(x, DeviceArray)