grgrjax

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.

Installation with pip

It's as simple as typing

pip install grgrjax

in the terminal (Unix/Mac) or the Anaconda Prompt (Win).