README

Transformable numerical computing at scale




Transformations
| Scaling
| Install guide
| Change logs
| Reference docs


 What is JAX?

JAX is a Python library for accelerator-oriented array computation and program transformation,
designed for high-performance numerical computing and large-scale machine learning.

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 jax.grad as well as forward-mode differentiation,
and the two can be composed arbitrarily to any order.

JAX uses XLA
to compile and scale your NumPy programs on TPUs, GPUs, and other hardware accelerators.
You can compile your own pure functions with jax.jit.
Compilation and automatic differentiation can be composed arbitrarily.

Dig a little deeper, and you'll see that JAX is really an extensible system for
composable function transformations at scale.

This is a research project, not an official Google product. Expect
sharp edges.
Please help by trying it out, reporting bugs,
and letting us know what you think!

python
import jax
import jax.numpy as jnp

def predict(params, inputs):
  for W, b in params:
    outputs = jnp.dot(inputs, W) + b
    inputs = jnp.tanh(outputs)   inputs to the next layer
  return outputs                 no activation on last layer

def loss(params, inputs, targets):
  preds = predict(params, inputs)
  return jnp.sum((preds - targets)2)

grad_loss = jax.jit(jax.grad(loss))   compiled gradient evaluation function
perex_grads = jax.jit(jax.vmap(grad_loss, in_axes=(None, 0, 0)))   fast per-example grads


 Contents
 Transformations
 Scaling
 Current gotchas
 Installation
 Citing JAX
 Reference documentation

 Transformations

At its core, JAX is an extensible system for transforming numerical functions.
Here are three: jax.grad, jax.jit, and jax.vmap.

 Automatic differentiation with grad

Use jax.grad
to efficiently compute reverse-mode gradients:

python
import jax
import jax.numpy as jnp

def tanh(x):
  y = jnp.exp(-2.0  x)
  return (1.0 - y) / (1.0 + y)

grad_tanh = jax.grad(tanh)
print(grad_tanh(1.0))
 prints 0.4199743


You can differentiate to any order with grad:

python
print(jax.grad(jax.grad(jax.grad(tanh)))(1.0))
 prints 0.62162673


You're free to use differentiation with Python control flow:

python
def abs_val(x):
  if x  0:
    return x
  else:
    return -x

abs_val_grad = jax.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 JAX Autodiff
Cookbook
and the reference docs on automatic
differentiation
for more.

 Compilation with jit

Use XLA to compile your functions end-to-end with
jit,
used either as an @jit decorator or as a higher-order function.

python
import jax
import jax.numpy as jnp

def slow_f(x):
   Element-wise ops see a large benefit from

Top issues

Discussions

Latest release

Bug fixes
   Built C++-generated intrinsic LLVM IR on Windows with clang-cl instead of
    emitting empty bitcode, enabling vectorized CPU intrinsics on Windows
    (https://github.com/openxla/xla/commit/3af77b5776fbc41a65d54876cc837bdc57540005).