The JAX Autodiff Cookbook

JAX’s autodiff is very general. It can calculate gradients of numpy functions, differentiating them with respect to nested lists, tuples and dicts. It can also calculate gradients of gradients and even work with complex numbers!

For more details, visit the source.