JAX API

Google's JAX autograd library is an example of a neat, composable API. The following is all that is needed to plot the gradient field of a function using JAX's grad and vmap functions.

Full page