Add JAX to the CPU/TPU image. (#992)

Fixes #918

Before, we were installing JAX only on the GPU image.
Included tests to prevent regression.

http://b/177334844

Author: rosbo

Dockerfile

@@ -420,6 +420,8 @@ RUN pip install flashtext && \
     # pycrypto is used by competitions team.
     pip install pycrypto && \
     pip install easyocr && \
+    # Keep JAX version in sync with GPU image.
+    pip install jax==0.2.12 jaxlib==0.1.64 && \
     /tmp/clean-layer.sh
 
 # Download base easyocr models.

gpu.Dockerfile

@@ -77,7 +77,7 @@ RUN pip uninstall -y lightgbm && \
     echo "libnvidia-opencl.so.1" > /etc/OpenCL/vendors/nvidia.icd && \
     /tmp/clean-layer.sh
 
-# Install JAX
+# Install JAX (Keep JAX version in sync with CPU image)
 RUN  pip install jax==0.2.12 jaxlib==0.1.64+cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION -f https://storage.googleapis.com/jax-releases/jax_releases.html && \
      /tmp/clean-layer.sh
 

tests/test_jax.py

@@ -1,22 +1,18 @@
 import unittest
-
 import time
 
+import jax.numpy as np
+
 from common import gpu_test
+from jax import grad, jit
 
 
 class TestJAX(unittest.TestCase):
     def tanh(self, x):
-        import jax.numpy as np
         y = np.exp(-2.0 * x)
         return (1.0 - y) / (1.0 + y)
 
-    @gpu_test
-    def test_JAX(self):
-        # importing inside the gpu-only test because these packages can't be
-        # imported on the CPU image since they are not present there.
-        from jax import grad, jit
-
+    def test_grad(self):
         grad_tanh = grad(self.tanh)
         ag = grad_tanh(1.0)
         self.assertEqual(0.4199743, ag)