This library provides some useful additional tools to JAX, such as a multivariate Newton method, val_and_jacfwd or vectorized JVP and VJP functions.
The module documentation can be found on Readthedocs.
It's as simple as typing
pip install grgrjaxin the terminal (Unix/Mac) or the Anaconda Prompt (Win).