feibyte
diff --git a/‎gpu.Dockerfile‎
Lines changed: 9 additions & 0 deletions b/‎gpu.Dockerfile‎
Lines changed: 9 additions & 0 deletions
diff --git a/‎tests/test_jax.py‎
Lines changed: 22 additions & 0 deletions b/‎tests/test_jax.py‎
Lines changed: 22 additions & 0 deletions
@@ -44,6 +44,15 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 && \
     /tmp/clean-layer.sh
 
+# Install JAX
+ENV JAX_PYTHON_VERSION=cp36
+ENV JAX_CUDA_VERSION=cuda100
+ENV JAX_PLATFORM=linux_x86_64
+ENV JAX_BASE_URL="https://storage.googleapis.com/jax-releases"
+
+RUN  pip install --upgrade $JAX_BASE_URL/$JAX_CUDA_VERSION/jaxlib-0.1.36-$JAX_PYTHON_VERSION-none-$JAX_PLATFORM.whl && \
+     pip install --upgrade jax
+
 # Reinstall packages with a separate version for GPU support.
 COPY --from=tensorflow_whl /tmp/tensorflow_gpu/*.whl /tmp/tensorflow_gpu/
 RUN pip uninstall -y tensorflow && \
 
@@ -0,0 +1,22 @@
+import unittest
+
+import time
+
+from common import gpu_test
+
+
+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
+
+        grad_tanh = grad(self.tanh)
+        ag = grad_tanh(1.0)
+        self.assertEqual(0.4199743, ag)