Unpin JAX (#1350)

A new version of `tensorflowjs` was released yesterday (4.16.0): https://pypi.org/project/tensorflowjs/#history

This include a fix for the incompatibility with the latest version of
JAX.

http://b/316967430

Author: rosbo

Dockerfile.tmpl

@@ -157,12 +157,11 @@ RUN pip install lightgbm==$LIGHTGBM_VERSION && \
 {{ end }}
 
 # Install JAX
-# b/316967430 Remove pin once new version of tensorflowjs is released (> 4.15.0)
 {{ if eq .Accelerator "gpu" }}
-RUN pip install "jax[cuda11_local]==0.4.21" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html && \
+RUN pip install "jax[cuda11_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html && \
     /tmp/clean-layer.sh
 {{ else }}
-RUN pip install jax[cpu]==0.4.21 && \
+RUN pip install jax[cpu] && \
     /tmp/clean-layer.sh
 {{ end }}