vmap is a JAX function transformation that defines the application of a function across one axis of an array.

This is an overview page of my notes on differential computing and JAX.

Contents:

1. What is differential computing
2. What automatic differentiation systems exist
3. Where differential computing gets used
4. Differential computing with JAX
5. What makes JAX special
6. Composable program transforms
7. Infinitely Wide Neural Networks

Stuff I've built/done with JAX:

1. jax-unirep
2. Workshop on Differential Learning
3. Estimating a multivariate Gaussian's parameters by gradient descent

lax.scan gives us loops with carryover.

Examples include:

• MCMC sampling
• gradient descent training routines
• recurrent neural networks

jit does Just In Time (JIT) compilation of a function to lower level code, thus bypassing the Python interpreter overhead.

grad returns a transformed version of a function f, which when called gives us a gradient function (let's call it df) gives us the derivative of f evaluated at the same arguments passed into f.