Autodiff is certainly one of the strengths of JAX. See the JAX Autodiff cookbook for a flavor of what JAX can do: https://colab.research.google.com/github/googl

Autodiff is certainly one of the strengths of JAX. See the JAX Autodiff cookbook for a flavor of what JAX can do: https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/autodiff_cookbook.ipynb https://colab.research.google.com/github/google/jax/blob/mas... You might also like the per-example gradients example that appears first in the JAX Github page: this is only one line of code, but important for research areas such as differential privacy.