JAX: Differentiable Computing by Google | by Branislav Holländer | Towards Data Science

Since deep learning took off in the early 2010s, many frameworks were written to facilitate deep learning in both research and production. For the record, let us mention Caffe, Theano, Torch, Lasagne, Tensorflow, Keras or PyTorch. Some of these frameworks disappeared after a while. Others survived and thrive to this day, mainly PyTorch and Tensorflow. Over time, these frameworks evolved into large ecosystems with many different functionalities. On one hand, they support the training of new models including massively parallel training on many-GPU systems. On the other hand, they allow users to deploy the trained models easily in the cloud or on mobile devices.
Unfortunately, this streamlining sometimes comes at the cost of flexibility required by machine learning researchers to test out new model ideas. It is precisely this user group that Google is targeting with its new framework called JAX. In this blog post, I will show you how JAX encourages experimentation by providing a low-level, high-performance interface to vectorized computing and gradient computation.
The goal of JAX is to allow the user to speed up raw Python and NumPy functions by just-in-time compilation and auto-parallelization as well as to compute gradients of these functions. In order to do this, JAX employs a functional design. Functions such as gradient computation are implemented as functional transforms (or functors) acting on user-defined Python functions. For instance, to calculate the gradient of the absolute value function, you can write:
from jax import graddef abs_val(x):  if x > 0:    return x  else:    return -xabs_val_grad = grad(abs_val)
As you can see, abs_val is a normal Python function, which is transformed by the functor grad.
In order to be able to use JAX functors, the user-defined functions have to follow some restrictions:
Every function processed by JAX is required to be pure. This means that it should always return the same result when called with the same input. In functional terms, it may not have any side effects. Otherwise, correctness when using just-in-time compilation or parallelization cannot be guaranteed by JAX. Functional purity isn’t actually enforced. Although errors may be thrown sometimes, it is mainly the responsibility of the programmer to ensure this condition.Instead of in-place mutating updates such as x[i] += 1, one has to use functional alternatives in the jax.ops package.Just-In-Time (JIT) compilation has some restrictions on Python control flow.
Every function processed by JAX is required to be pure. This means that it should always return the same result when called with the same input. In functional terms, it may not have any side effects. Otherwise, correctness when using just-in-time compilation or parallelization cannot be guaranteed by JAX. Functional purity isn’t actually enforced. Although errors may be thrown sometimes, it is mainly the responsibility of the programmer to ensure this condition.
Instead of in-place mutating updates such as x[i] += 1, one has to use functional alternatives in the jax.ops package.
Just-In-Time (JIT) compilation has some restrictions on Python control flow.
In terms of philosophy, JAX is similar to what you would experience when using purely functional languages such as Haskell. In fact, the developers even use Haskell-type signatures in the documentation to explain some of the transformations.
One of the main features of JAX is the ability to speed up execution of Python code by JIT. Internally, JAX uses the XLA compiler to accomplish this. XLA is able to compile code not only for CPUs, but also for GPUs or even TPUs. This makes JAX very powerful and versatile. In order to compile a function using XLA, you can use the jit functor like this:
from jax import jitimport jax.numpy as jnpdef selu(x, alpha=1.67, lmbda=1.05):  return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)selu_jit = jit(selu)
Alternatively, you may also use jit as a decorator on top of the function definition:
@jitdef selu(x, alpha=1.67, lmbda=1.05):  return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)
The jit functor transforms the input function by returning a compiled version of it to the caller. Calling the original selu function will use the Python interpreter, while calling selu_jit will call the compiled version, which should be much faster, especially for vectorized inputs such as numpy arrays. Furthermore, the JIT compilation only occurs once and is cached thereafter, making subsequent calls of the function very efficient.
When training machine learning models, you are usually computing some loss for a subset of input data, then updating the model parameters. Since computing the forward function of the model for every input sequentially would be too slow, it is customary to batch together the subset of data into one batch and computing the forward function in a parallel fashion on an accelerator such as the GPU. This is accomplished in JAX by using the function vmap:
def forward(input, w, b):    input = jnp.dot(input, w) + b    return jnp.max(input, 0)jax.vmap(forward, in_axes=(0, None, None), out_axes=0)(inputs, w, b)
This code snippet takes a forward function of a fully-connected layer with ReLU activation and parallelizes its execution over a batch of inputs. The parameters in_axes and out_axes specify over which parameters and axes the parallelization occurs. In this case, (0, None, None) means that the input is parallelized over the 0-th axis while w and b are left untouched. The output is parallelized over the 0-th axis.
What makes JAX particularly interesting to machine learning researchers is its ability to compute gradients of arbitrary pure functions. JAX inherited this capability from autograd, a package to compute derivatives of NumPy arrays.
To compute the gradient of a function, simply use the grad transformation:
import jax.numpy as jnpgrad(jnp.tanh))(2.0)[0.070650816]
If you want to compute higher-order derivatives, you can simply chain together multiple grad transformations like this:
import jax.numpy as jnpgrad(grad(jnp.tanh))(2.0)[-0.13621889]
While applying grad to a function in R (the set of real numbers) will give you a single number as the output, you may also apply it to a vector-valued function to obtain a Jacobian:
def f(x):    return jnp.asarray(        [x[0], 5*x[2], 4*x[1]**2 - 2*x[2], x[2] * jnp.sin(x[0])])print(jax.jacfwd(f)(jnp.array([1., 2., 3.])))[[ 1.       0.       0.     ] [ 0.       0.       5.     ] [ 0.      16.      -2.     ] [ 1.6209   0.       0.84147]]
Here, we do not use the grad transformation, since it only works on scalar-output functions. Instead, we use a similar function called jacfwd for automatic forward-mode differentiation (for a detailed discussion of forward- and backward-mode differentiation please refer to the JAX docs).
Note that we may chain together the functional transformations in JAX. For instance, we may vectorize our function using vmap, calculate the gradients using grad and then compile the resulting function using jit, such as in this example of a standard mini-batch training loop in deep learning:
def loss(x, y):    out = net(x)    cross_entropy = -y * np.log(out) - (1 - y)*np.log(1 - out)    return cross_entropyloss_grad = jax.jit(jax.vmap(jax.grad(loss), in_axes=(0, 0), out_axes=0))
Here, loss_grad computes the gradient of the cross-entropy loss, while paralleling over the input (x, y) and the single output (the loss). The entire function is jitted to allow for fast computation on the GPU. The resulting function is cached, allowing it to be called with different outputs without any further overhead.
The composition of multiple transformations makes the framework extremely powerful and gives the user great flexibility in designing the data flow.
JAX provides a useful alternative to more high-level frameworks such as PyTorch or Tensorflow for researchers who need extra flexibility. The ability to differentiate through native Python and NumPy functions is amazing and the JIT compilation and auto-vectorization features greatly simplify writing efficient code for massively-parallel architectures like GPUs or TPUs. What I like most about JAX, however, is its clean functional interface. Surely it’s only a matter of time before an ecosystem of mature high-level APIs will evolve around it. Stay tuned!