## JAX: Autograd and XLA

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more.

## What is JAX?

JAX is Autograd and

XLA,

brought together for high-performance machine learning research.

With its updated version of Autograd,

JAX can automatically differentiate native

Python and NumPy functions. It can differentiate through loops, branches,

recursion, and closures, and it can take derivatives of derivatives of

derivatives. It supports reverse-mode differentiation (a.k.a. backpropagation)

via `grad`

as well as forward-mode differentiation,

and the two can be composed arbitrarily to any order.

What’s new is that JAX uses

XLA

to compile and run your NumPy programs on GPUs and TPUs. Compilation happens

under the hood by default, with library calls getting just-in-time compiled and

executed. But JAX also lets you just-in-time compile your own Python functions

into XLA-optimized kernels using a one-function API,

`jit`

. Compilation and automatic differentiation can be

composed arbitrarily, so you can express sophisticated algorithms and get

maximal performance without leaving Python. You can even program multiple GPUs

or TPU cores at once using `pmap`

, and

differentiate through the whole thing.

Dig a little deeper, and you'll see that JAX is really an extensible system for

composable function transformations. Both

`grad`

and `jit`

are instances of such transformations. Others are

`vmap`

for automatic vectorization and

`pmap`

for single-program multiple-data (SPMD)

parallel programming of multiple accelerators, with more to come.

This is a research project, not an official Google product. Expect bugs and

sharp edges.

Please help by trying it out, reporting

bugs, and letting us know what you

think!

```
import jax.numpy as jnp
from jax import grad, jit, vmap
def predict(params, inputs):
for W, b in params:
outputs = jnp.dot(inputs, W) + b
inputs = jnp.tanh(outputs)
return outputs
def logprob_fun(params, inputs, targets):
preds = predict(params, inputs)
return jnp.sum((preds - targets)**2)
grad_fun = jit(grad(logprob_fun)) # compiled gradient evaluation function
perex_grads = jit(vmap(grad_fun, in_axes=(None, 0, 0))) # fast per-example grads
```

## Quickstart: Colab in the Cloud

Jump right in using a notebook in your browser, connected to a Google Cloud GPU.

Here are some starter notebooks:

- The basics: NumPy on accelerators,
`grad`

for differentiation,`jit`

for compilation, and`vmap`

for vectorization - Training a Simple Neural Network, with TensorFlow Dataset Data Loading

**JAX now runs on Cloud TPUs.** To try out the preview, see the Cloud TPU

Colabs.

For a deeper dive into JAX:

- The Autodiff Cookbook, Part 1: easy and powerful automatic differentiation in JAX
- Common gotchas and sharp edges
- See the full list of

notebooks.

You can also take a look at the mini-libraries in

`jax.experimental`

,

like `stax`

for building neural

networks

and `optimizers`

for first-order stochastic

optimization,

or the examples.

## Transformations

At its core, JAX is an extensible system for transforming numerical functions.

Here are four of primary interest: `grad`

, `jit`

, `vmap`

, and `pmap`

.

### Automatic differentiation with `grad`

JAX has roughly the same API as Autograd.

The most popular function is

`grad`

for reverse-mode gradients:

```
from jax import grad
import jax.numpy as jnp
def tanh(x): # Define a function
y = jnp.exp(-2.0 * x)
return (1.0 - y) / (1.0 + y)
grad_tanh = grad(tanh) # Obtain its gradient function
print(grad_tanh(1.0)) # Evaluate it at x = 1.0
# prints 0.4199743
```

You can differentiate to any order with `grad`

.

```
print(grad(grad(grad(tanh)))(1.0))
# prints 0.62162673
```

For more advanced autodiff, you can use

`jax.vjp`

for

reverse-mode vector-Jacobian products and

`jax.jvp`

for

forward-mode Jacobian-vector products. The two can be composed arbitrarily with

one another, and with other JAX transformations. Here's one way to compose those

to make a function that efficiently computes full Hessian

matrices:

```
from jax import jit, jacfwd, jacrev
def hessian(fun):
return jit(jacfwd(jacrev(fun)))
```

As with Autograd, you're free to use

differentiation with Python control structures:

```
def abs_val(x):
if x > 0:
return x
else:
return -x
abs_val_grad = grad(abs_val)
print(abs_val_grad(1.0)) # prints 1.0
print(abs_val_grad(-1.0)) # prints -1.0 (abs_val is re-evaluated)
```

See the reference docs on automatic

differentiation

and the JAX Autodiff

Cookbook

for more.

### Compilation with `jit`

You can use XLA to compile your functions end-to-end with

`jit`

,

used either as an `@jit`

decorator or as a higher-order function.

```
import jax.numpy as jnp
from jax import jit
def slow_f(x):
# Element-wise ops see a large benefit from fusion
return x * x + x * 2.0
x = jnp.ones((5000, 5000))
fast_f = jit(slow_f)
%timeit -n10 -r3 fast_f(x) # ~ 4.5 ms / loop on Titan X
%timeit -n10 -r3 slow_f(x) # ~ 14.5 ms / loop (also on GPU via JAX)
```

You can mix `jit`

and `grad`

and any other JAX transformation however you like.

Using `jit`

puts constraints on the kind of Python control flow

the function can use; see

the Gotchas

Notebook

for more.

### Auto-vectorization with `vmap`

`vmap`

is

the vectorizing map.

It has the familiar semantics of mapping a function along array axes, but

instead of keeping the loop on the outside, it pushes the loop down into a

function’s primitive operations for better performance.

Using `vmap`

can save you from having to carry around batch dimensions in your

code. For example, consider this simple *unbatched* neural network prediction

function:

```
def predict(params, input_vec):
assert input_vec.ndim == 1
for W, b in params:
output_vec = jnp.dot(W, input_vec) + b # `input_vec` on the right-hand side!
input_vec = jnp.tanh(output_vec)
return output_vec
```

We often instead write `jnp.dot(inputs, W)`

to allow for a batch dimension on the

left side of `inputs`

, but we’ve written this particular prediction function to

apply only to single input vectors. If we wanted to apply this function to a

batch of inputs at once, semantically we could just write

```
from functools import partial
predictions = jnp.stack(list(map(partial(predict, params), input_batch)))
```

But pushing one example through the network at a time would be slow! It’s better

to vectorize the computation, so that at every layer we’re doing matrix-matrix

multiplication rather than matrix-vector multiplication.

The `vmap`

function does that transformation for us. That is, if we write

```
from jax import vmap
predictions = vmap(partial(predict, params))(input_batch)
# or, alternatively
predictions = vmap(predict, in_axes=(None, 0))(params, input_batch)
```

then the `vmap`

function will push the outer loop inside the function, and our

machine will end up executing matrix-matrix multiplications exactly as if we’d

done the batching by hand.

It’s easy enough to manually batch a simple neural network without `vmap`

, but

in other cases manual vectorization can be impractical or impossible. Take the

problem of efficiently computing per-example gradients: that is, for a fixed set

of parameters, we want to compute the gradient of our loss function evaluated

separately at each example in a batch. With `vmap`

, it’s easy:

```
per_example_gradients = vmap(partial(grad(loss), params))(inputs, targets)
```

Of course, `vmap`

can be arbitrarily composed with `jit`

, `grad`

, and any other

JAX transformation! We use `vmap`

with both forward- and reverse-mode automatic

differentiation for fast Jacobian and Hessian matrix calculations in

`jax.jacfwd`

, `jax.jacrev`

, and `jax.hessian`

.

### SPMD programming with `pmap`

For parallel programming of multiple accelerators, like multiple GPUs, use

`pmap`

.

With `pmap`

you write single-program multiple-data (SPMD) programs, including

fast parallel collective communication operations. Applying `pmap`

will mean

that the function you write is compiled by XLA (similarly to `jit`

), then

replicated and executed in parallel across devices.

Here's an example on an 8-GPU machine:

```
from jax import random, pmap
import jax.numpy as jnp
# Create 8 random 5000 x 6000 matrices, one per GPU
keys = random.split(random.PRNGKey(0), 8)
mats = pmap(lambda key: random.normal(key, (5000, 6000)))(keys)
# Run a local matmul on each device in parallel (no data transfer)
result = pmap(lambda x: jnp.dot(x, x.T))(mats) # result.shape is (8, 5000, 5000)
# Compute the mean on each device in parallel and print the result
print(pmap(jnp.mean)(result))
# prints [1.1566595 1.1805978 ... 1.2321935 1.2015157]
```

In addition to expressing pure maps, you can use fast collective communication

operations

between devices:

```
from functools import partial
from jax import lax
@partial(pmap, axis_name='i')
def normalize(x):
return x / lax.psum(x, 'i')
print(normalize(jnp.arange(4.)))
# prints [0. 0.16666667 0.33333334 0.5 ]
```

You can even nest `pmap`

functions for more

sophisticated communication patterns.

It all composes, so you're free to differentiate through parallel computations:

```
from jax import grad
@pmap
def f(x):
y = jnp.sin(x)
@pmap
def g(z):
return jnp.cos(z) * jnp.tan(y.sum()) * jnp.tanh(x).sum()
return grad(lambda w: jnp.sum(g(w)))(x)
print(f(x))
# [[ 0. , -0.7170853 ],
# [-3.1085174 , -0.4824318 ],
# [10.366636 , 13.135289 ],
# [ 0.22163185, -0.52112055]]
print(grad(lambda x: jnp.sum(f(x)))(x))
# [[ -3.2369726, -1.6356447],
# [ 4.7572474, 11.606951 ],
# [-98.524414 , 42.76499 ],
# [ -1.6007166, -1.2568436]]
```

When reverse-mode differentiating a `pmap`

function (e.g. with `grad`

), the

backward pass of the computation is parallelized just like the forward pass.

See the SPMD

Cookbook

and the SPMD MNIST classifier from scratch

example

for more.

## Current gotchas

For a more thorough survey of current gotchas, with examples and explanations,

we highly recommend reading the Gotchas

Notebook.

Some standouts:

- JAX transformations only work on pure functions, which don't have side-effects and respect referential transparency (i.e. object identity testing with
`is`

isn't preserved). If you use a JAX transformation on an impure Python function, you might see an error like`Exception: Can't lift Traced...`

or`Exception: Different traces at same level`

. - In-place mutating updates of

arrays, like`x[i] += y`

, aren't supported, but there are functional alternatives. Under a`jit`

, those functional alternatives will reuse buffers in-place automatically. - Random numbers are

different, but for good reasons. - If you're looking for convolution

operators,

they're in the`jax.lax`

package. - JAX enforces single-precision (32-bit, e.g.
`float32`

) values by default, and

to enable

double-precision

(64-bit, e.g.`float64`

) one needs to set the`jax_enable_x64`

variable at

startup (or set the environment variable`JAX_ENABLE_X64=True`

). - Some of NumPy's dtype promotion semantics involving a mix of Python scalars

and NumPy types aren't preserved, namely`np.add(1, np.array([2], np.float32)).dtype`

is`float64`

rather than`float32`

. - Some transformations, like
`jit`

, constrain how you can use Python control

flow.

You'll always get loud errors if something goes wrong. You might have to use

`jit`

's`static_argnums`

parameter,

structured control flow

primitives

like

`lax.scan`

,

or just use`jit`

on smaller subfunctions.

## Installation

JAX is written in pure Python, but it depends on XLA, which needs to be

installed as the `jaxlib`

package. Use the following instructions to install a

binary package with `pip`

, or to build JAX from source.

We support installing or building `jaxlib`

on Linux (Ubuntu 16.04 or later) and

macOS (10.12 or later) platforms. Windows users can use JAX on CPU and GPU via

the

Windows Subsystem for Linux.

There is some initial native Windows support, but since it is still somewhat

immature, there are no binary releases and it must be

built from source.

### pip installation

To install a CPU-only version, which might be useful for doing local

development on a laptop, you can run

```
pip install --upgrade pip
pip install --upgrade jax jaxlib # CPU-only version
```

On Linux, it is often necessary to first update `pip`

to a version that supports

`manylinux2010`

wheels.

If you want to install JAX with both CPU and GPU support, using existing CUDA

and CUDNN7 installations on your machine (for example, preinstalled on your

cloud VM), you can run

```
pip install --upgrade pip
pip install --upgrade jax jaxlib==0.1.57+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html
```

The jaxlib version must correspond to the version of the existing CUDA

installation you want to use, with `cuda110`

for CUDA 11.0, `cuda102`

for CUDA

10.2, and `cuda101`

for CUDA 10.1. You can find your

CUDA version with: install path:

```
nvcc --version
```

Note that some GPU functionality expects the CUDA installation to be at

`/usr/local/cuda-X.X`

, where X.X should be replaced with the CUDA version number

(e.g. `cuda-10.2`

). If CUDA is installed elsewhere on your system, you can either

create a symlink:

```
sudo ln -s /path/to/cuda /usr/local/cuda-X.X
```

Or set the following environment variable before importing JAX:

```
XLA_FLAGS=--xla_gpu_cuda_data_dir=/path/to/cuda
```

Please let us know on the issue tracker

if you run into any errors or problems with the prebuilt wheels.

### Building JAX from source

## Neural network libraries

Multiple Google research groups develop and share libraries for training neural

networks in JAX. If you want a fully featured library for neural network

training with examples and how-to guides, try

Flax. Another option is

Trax, a combinator-based framework focused on

ease-of-use and end-to-end single-command examples, especially for sequence

models and reinforcement learning. Finally,

Objax is a minimalist object-oriented

framework with a PyTorch-like interface.

DeepMind has open-sourced an ecosystem of libraries around JAX including

Haiku for neural network modules,

Optax for gradient processing and

optimization, RLax for RL algorithms, and

chex for reliable code and testing.

## Citing JAX

To cite this repository:

```
@software{jax2018github,
author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang},
title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs},
url = {http://github.com/google/jax},
version = {0.2.5},
year = {2018},
}
```

In the above bibtex entry, names are in alphabetical order, the version number

is intended to be that from jax/version.py, and

the year corresponds to the project's open-source release.

A nascent version of JAX, supporting only automatic differentiation and

compilation to XLA, was described in a paper that appeared at SysML

2018. We're currently working on

covering JAX's ideas and capabilities in a more comprehensive and up-to-date

paper.

## Reference documentation

For details about the JAX API, see the

reference documentation.

For getting started as a JAX developer, see the

developer documentation.