jax

f538df4d - Replace the docker image with pre-installed CUDA with a clean Docker image.