more things

Dockerfile

@@ -6,7 +6,7 @@ RUN \
   apt-get -qq -y upgrade && \
   apt-get install -y jq git tree hdf5-tools bash-completion && \
   apt-get install -y ffmpeg vim emacs && \
-  apt-get install -y gcc libopenmpi-dev openmpi-bin && \
+  apt-get install -y gcc libopenmpi-dev openmpi-bin graphviz && \
   apt-get clean && \
   rm -rf /var/lib/apt-get/lists/* && \
   true
@@ -26,3 +26,4 @@ USER grk
 WORKDIR /home/grk
 
 CMD /bin/bash
+

requirements.txt

@@ -20,3 +20,6 @@ numba
 mpi4py
 Cython
 sympy
+jax
+jaxlib
+graphviz