You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
There is some demand for a Dockerfile that runs jax / jaxlib (as opposed to the one currently in the repository that is used to build wheels for jax/jaxlib), see here, here, and probably here.
The Dockerfile
Building a dockerfile for jax usage as opposed to building jax is probably low on the list of priorities, but I've found empirically that jax "just works" with tensorflow-gpu images. An example dockerfile that takes a tensorflow-gpu image (2.4.0 in this case) and adds on jax/jaxlib via pip-install is given below. This is not an official JAX image! No guarantees it will continue working
FROM tensorflow/tensorflow:2.4.0-gpu
# this contains CUDA 11.0 and CUDNN 8.X.X
# you probably want to create a new user as an entry point, not run root inside the container
RUN useradd -ms /bin/bash jax
# switch into the jax user
USER jax
# add the user's .local/bin to path so that our --user pip installed packages will appear
ENV PATH=/home/jax/.local/bin:$PATH
# install python packages via pip
# for CUDA support, we have to specify a particular jaxlib version
# we also need to specify the particular cuda version
RUN pip3 install --user \
jax==0.2.10 \
jaxlib==0.1.62+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html
CMD /bin/bash
For those trying to avoid tensorflow bloat, this issue has a suggested Dockerfile for building your own jaxlib on nvidia's cuda cudnn devel images.
reacted with thumbs up emoji reacted with thumbs down emoji reacted with laugh emoji reacted with hooray emoji reacted with confused emoji reacted with heart emoji reacted with rocket emoji reacted with eyes emoji
-
Demand for JAX Dockerfile
There is some demand for a Dockerfile that runs jax / jaxlib (as opposed to the one currently in the repository that is used to build wheels for jax/jaxlib), see here, here, and probably here.
The Dockerfile
Building a dockerfile for jax usage as opposed to building jax is probably low on the list of priorities, but I've found empirically that jax "just works" with tensorflow-gpu images. An example dockerfile that takes a tensorflow-gpu image (2.4.0 in this case) and adds on jax/jaxlib via pip-install is given below. This is not an official JAX image! No guarantees it will continue working
For those trying to avoid tensorflow bloat, this issue has a suggested Dockerfile for building your own jaxlib on nvidia's cuda cudnn devel images.
Beta Was this translation helpful? Give feedback.
All reactions