Import tpu_driver after xla_client (jax-ml#3064)

This is a workaround until we build a new jaxlib with tensorflow/tensorflow@f462867
BuddenD · May 12, 2020 · 0d97c3b · 0d97c3b
1 parent 8008aa9
commit 0d97c3b
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/jax/lib/__init__.py b/jax/lib/__init__.py
@@ -45,10 +45,6 @@ def _check_jaxlib_version():
 _check_jaxlib_version()
 
 
-try:
-  from jaxlib import tpu_client  # pytype: disable=import-error
-except:
-  tpu_client = None
 from jaxlib import xla_client
 from jaxlib import lapack
 
@@ -58,3 +54,8 @@ def _check_jaxlib_version():
   from jaxlib import cuda_prng
 except ImportError:
   cuda_prng = None
+
+try:
+  from jaxlib import tpu_client  # pytype: disable=import-error
+except:
+  tpu_client = None