Jax is not compiled for GPU

[Selimonder]

🐛 Bug

Expecting to see a GPU device on backend, but seeing CPU instead.

This wasn't an issue on earlier environments. Also, thanks for making Flax a default package 🙌

[Selimonder]

temporary fix for offline environment (latest as 14th of July): use output data of https://www.kaggle.com/realsleim/jax-wheels

execute following in the beginning of your notebook

!pip install -U ../input/jax-wheels/pip-22.1.2-py3-none-any.whl --no-index --find-links ../input/jax-wheels/
!pip install ../input/jax-wheels/jaxlib-0.3.14+cuda11.cudnn82-cp37-none-manylinux2014_x86_64.whl --no-index --find-links ../input/jax-wheels/

result

[rosbo]

Hi @Selimonder,

Looks like this installs a JAX 2.x with GPU support and later in the build, the FLAX install upgrade JAX to 3.x with CPU only support:

docker-python/Dockerfile.tmpl

Lines 126 to 127 in 029dea1

	RUN pip install jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_releases.html && \
	/tmp/clean-layer.sh

I will fix this and add a test to prevent regression on this.

Thanks for flagging.

[rosbo]

I found the issue: the CUDA releases have been moved to: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html (before they were available at https://storage.googleapis.com/jax-releases/jax_releases.html).

[rosbo]

@Selimonder A new image with the fix will be released tomorrow.