Add JAX dependencies to Dockerfile. (Note at time of PR this isn't super useful as the base image is not GPU compatible, but can still be used for autodiff).
Also use python -m pip --no-cache-dir option for best practices in Docker and to reduce the Docker image size.
* Add JAX dependencies to Dockerfile to use autodiff
* Use pip --no-cache-dir option for best practices
* Add base image pull to Makefile
Add JAX dependencies to Dockerfile. (Note at time of PR this isn't super useful as the base image is not GPU compatible, but can still be used for autodiff).
Also use
python -m pip --no-cache-dir
option for best practices in Docker and to reduce the Docker image size.