JAX's grad function is merely the "gateway drug" to this bigger idea of "composable program transforms". grad is one example of a composable program transform: that is transforming one program into another in a composable fashion. Other transforms include vmap, lax.scan, jit, and more. These all accept Python functions and return Python functions. jit, in particular, can accelerate a program anywhere from 2-100 fold on a CPU, depending on what your reasonable baseline comparison is.

In particular, the latter three of the aformentioned transforms allow for highly performant loop code, written without loops, that can also be composed together. In a differential learning workshop that I have been developing, I provide further details in there, which you can take a look at.

There are other automatic program transformations that are in active development, and one exciting realm I see is in the probabilistic programming world. In PyMC3, for example, an automated transform happens when we take our PyMC3 syntax, which is written in a domain specific language (DSL) implemented in Python, and transform/compile it into a compute graph that gives us a likelihood function. It's as if PyMC3 gives us a likelihood(func) analogy to JAX's grad func. (If you've tried writing probabilistic model likelihoods by hand, you'll know how much of a convenience this is!)