Building JAX on the Auton cluster

[Viraj Mehta]

Hi everyone,

I had a bit of an adventure trying to build JAX on Auton and thought I’d document the right way to do it so that if others want they can as well. Here are the steps to do it on a machine that has CUDA 11 and cuDNN 8:

1. Go somewhere in your scratch
2. git clone https://github.com/google/jax
3. Make a python environment with a python >= 3.6 using Conda or virtualenv
4. Install numpy, spicy, cython, six
5. Source /opt/rh/devtoolset-8/enable
6. Run python build/build.py —enable_cuda —cuda_path /usr/local/cuda —cudnn_path /usr/
7. Pip install -e build
8. Pip install -e .

I know this isn’t that complicated, but I figured it would save some effort if anyone else would like to use it.

Cheers,
Viraj
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mailman.srv.cs.cmu.edu/pipermail/autonlab-users/attachments/20200831/e49dff4b/attachment-0001.html>