Code doesn't work with newest jax, jaxlib & neural-tangents libraries

Unfortunately the code base doesn't work with the newest versions of those three libraries. Ideally we would provide the user with an easy way to install the correct versions such that they can run the code as-is.