Upgrade JAX to 0.2.16 (#1039)

Author: rosbo

gpu.Dockerfile

@@ -78,8 +78,8 @@ RUN pip uninstall -y lightgbm && \
     /tmp/clean-layer.sh
 
 # 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
+RUN pip install jax==0.2.16 jaxlib==0.1.68+cuda$CUDA_MAJOR_VERSION$CUDA_MINOR_VERSION -f https://storage.googleapis.com/jax-releases/jax_releases.html && \
+    /tmp/clean-layer.sh
 
 # Reinstall packages with a separate version for GPU support.
 RUN pip uninstall -y tensorflow && \