Package: jax
JAX: Autograd and XLAQuickstart|Transformations|Install guide|Neural net libraries|Change logs|Reference docsWhat is JAX?JAX isAutogradandXLA,
brought together for high-performance numerical computing, including
large-scale machine learning research.With its updated version ofAutograd,
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)
viagradas well as forward-mode differentiation,
and the two can be composed arbitrarily to any order.What’s new is that JAX usesXLAto 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 usingpmap, and
differentiate through the whole thing.Dig a little deeper, and you'll see that JAX is really an extensible system forcomposable function transformations. Bothgradandjitare instances of such transformations. Others arevmapfor automatic vectorization andpmapfor 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 andsharp edges.
Please help by trying it out,reporting
bugs, and letting us know what you
think!importjax.numpyasjnpfromjaximportgrad,jit,vmapdefpredict(params,inputs):forW,binparams:outputs=jnp.dot(inputs,W)+binputs=jnp.tanh(outputs)# inputs to the next layerreturnoutputs# no activation on last layerdefloss(params,inputs,targets):preds=predict(params,inputs)returnjnp.sum((preds-targets)**2)grad_loss=jit(grad(loss))# compiled gradient evaluation functionperex_grads=jit(vmap(grad_loss,in_axes=(None,0,0)))# fast per-example gradsContentsQuickstart: Colab in the CloudTransformationsCurrent gotchasInstallationNeural net librariesCiting JAXReference documentationQuickstart: Colab in the CloudJump right in using a notebook in your browser, connected to a Google Cloud GPU.
Here are some starter notebooks:The basics: NumPy on accelerators,gradfor differentiation,jitfor compilation, andvmapfor vectorizationTraining a Simple Neural Network, with TensorFlow Dataset Data LoadingJAX now runs on Cloud TPUs.To try out the preview, see theCloud TPU
Colabs.For a deeper dive into JAX:The Autodiff Cookbook, Part 1: easy and powerful automatic differentiation in JAXCommon gotchas and sharp edgesSee thefull list of
notebooks.TransformationsAt its core, JAX is an extensible system for transforming numerical functions.
Here are four transformations of primary interest:grad,jit,vmap, andpmap.Automatic differentiation withgradJAX has roughly the same API asAutograd.
The most popular function isgradfor reverse-mode gradients:fromjaximportgradimportjax.numpyasjnpdeftanh(x):# Define a functiony=jnp.exp(-2.0*x)return(1.0-y)/(1.0+y)grad_tanh=grad(tanh)# Obtain its gradient functionprint(grad_tanh(1.0))# Evaluate it at x = 1.0# prints 0.4199743You can differentiate to any order withgrad.print(grad(grad(grad(tanh)))(1.0))# prints 0.62162673For more advanced autodiff, you can usejax.vjpfor
reverse-mode vector-Jacobian products andjax.jvpfor
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 computesfull Hessian
matrices:fromjaximportjit,jacfwd,jacrevdefhessian(fun):returnjit(jacfwd(jacrev(fun)))As withAutograd, you're free to use
differentiation with Python control structures:defabs_val(x):ifx>0:returnxelse:return-xabs_val_grad=grad(abs_val)print(abs_val_grad(1.0))# prints 1.0print(abs_val_grad(-1.0))# prints -1.0 (abs_val is re-evaluated)See thereference docs on automatic
differentiationand theJAX Autodiff
Cookbookfor more.Compilation withjitYou can use XLA to compile your functions end-to-end withjit,
used either as an@jitdecorator or as a higher-order function.importjax.numpyasjnpfromjaximportjitdefslow_f(x):# Element-wise ops see a large benefit from fusionreturnx*x+x*2.0x=jnp.ones((5000,5000))fast_f=jit(slow_f)%timeit-n10-r3fast_f(x)# ~ 4.5 ms / loop on Titan X%timeit-n10-r3slow_f(x)# ~ 14.5 ms / loop (also on GPU via JAX)You can mixjitandgradand any other JAX transformation however you like.Usingjitputs constraints on the kind of Python control flow
the function can use; see
theGotchas
Notebookfor more.Auto-vectorization withvmapvmapis
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.Usingvmapcan save you from having to carry around batch dimensions in your
code. For example, consider this simpleunbatchedneural network prediction
function:defpredict(params,input_vec):assertinput_vec.ndim==1activations=input_vecforW,binparams:outputs=jnp.dot(W,activations)+b# `activations` on the right-hand side!activations=jnp.tanh(outputs)# inputs to the next layerreturnoutputs# no activation on last layerWe often instead writejnp.dot(activations, W)to allow for a batch dimension on the
left side ofactivations, 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 writefromfunctoolsimportpartialpredictions=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.Thevmapfunction does that transformation for us. That is, if we writefromjaximportvmappredictions=vmap(partial(predict,params))(input_batch)# or, alternativelypredictions=vmap(predict,in_axes=(None,0))(params,input_batch)then thevmapfunction 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 withoutvmap, 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. Withvmap, it’s easy:per_example_gradients=vmap(partial(grad(loss),params))(inputs,targets)Of course,vmapcan be arbitrarily composed withjit,grad, and any other
JAX transformation! We usevmapwith both forward- and reverse-mode automatic
differentiation for fast Jacobian and Hessian matrix calculations injax.jacfwd,jax.jacrev, andjax.hessian.SPMD programming withpmapFor parallel programming of multiple accelerators, like multiple GPUs, usepmap.
Withpmapyou write single-program multiple-data (SPMD) programs, including
fast parallel collective communication operations. Applyingpmapwill mean
that the function you write is compiled by XLA (similarly tojit), then
replicated and executed in parallel across devices.Here's an example on an 8-GPU machine:fromjaximportrandom,pmapimportjax.numpyasjnp# Create 8 random 5000 x 6000 matrices, one per GPUkeys=random.split(random.PRNGKey(0),8)mats=pmap(lambdakey:random.normal(key,(5000,6000)))(keys)# Run a local matmul on each device in parallel (no data transfer)result=pmap(lambdax:jnp.dot(x,x.T))(mats)# result.shape is (8, 5000, 5000)# Compute the mean on each device in parallel and print the resultprint(pmap(jnp.mean)(result))# prints [1.1566595 1.1805978 ... 1.2321935 1.2015157]In addition to expressing pure maps, you can use fastcollective communication
operationsbetween devices:fromfunctoolsimportpartialfromjaximportlax@partial(pmap,axis_name='i')defnormalize(x):returnx/lax.psum(x,'i')print(normalize(jnp.arange(4.)))# prints [0.         0.16666667 0.33333334 0.5       ]You can evennestpmapfunctionsfor more
sophisticated communication patterns.It all composes, so you're free to differentiate through parallel computations:fromjaximportgrad@pmapdeff(x):y=jnp.sin(x)@pmapdefg(z):returnjnp.cos(z)*jnp.tan(y.sum())*jnp.tanh(x).sum()returngrad(lambdaw: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(lambdax: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 apmapfunction (e.g. withgrad), the
backward pass of the computation is parallelized just like the forward pass.See theSPMD
Cookbookand theSPMD MNIST classifier from scratch
examplefor more.Current gotchasFor a more thorough survey of current gotchas, with examples and explanations,
we highly recommend reading theGotchas
Notebook.
Some standouts:JAX transformations only work onpure functions, which don't have side-effects and respectreferential transparency(i.e. object identity testing withisisn't preserved). If you use a JAX transformation on an impure Python function, you might see an error likeException: Can't lift Traced...orException: Different traces at same level.In-place mutating updates of
arrays, likex[i] += y, aren't supported, butthere are functional alternatives. Under ajit, those functional alternatives will reuse buffers in-place automatically.Random numbers are
different, but forgood reasons.If you're looking forconvolution
operators,
they're in thejax.laxpackage.JAX enforces single-precision (32-bit, e.g.float32) values by default, andto enable
double-precision(64-bit, e.g.float64) one needs to set thejax_enable_x64variable at
startup (or set the environment variableJAX_ENABLE_X64=True).
On TPU, JAX uses 32-bit values by default for everythingexceptinternal
temporary variables in 'matmul-like' operations, such asjax.numpy.dotandlax.conv.
Those ops have aprecisionparameter which can be used to simulate
true 32-bit, with a cost of possibly slower runtime.Some of NumPy's dtype promotion semantics involving a mix of Python scalars
and NumPy types aren't preserved, namelynp.add(1, np.array([2], np.float32)).dtypeisfloat64rather thanfloat32.Some transformations, likejit,constrain how you can use Python control
flow.
You'll always get loud errors if something goes wrong. You might have to usejit'sstatic_argnumsparameter,structured control flow
primitiveslikelax.scan,
or just usejiton smaller subfunctions.InstallationSupported platformsLinux x86_64Linux aarch64Mac x86_64Mac ARMWindows x86_64Windows WSL2 x86_64CPUyesyesyesyesyesyesNVIDIA GPUyesyesnon/anoexperimentalGoogle TPUyesn/an/an/an/an/aAMD GPUexperimentalnonon/anonoApple GPUn/anoexperimentalexperimentaln/an/aInstructionsHardwareInstructionsCPUpip install -U "jax[cpu]"NVIDIA GPU on x86_64pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlGoogle TPUpip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.htmlAMD GPUUseDockerorbuild from source.Apple GPUFollowApple's instructions.Seethe documentationfor information on alternative installation strategies. These include compiling
from source, installing with Docker, using other versions of CUDA, a
community-supported conda build, and answers to some frequently-asked questions.Neural network librariesMultiple 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, tryFlax.Google X maintains the neural network libraryEquinox. This is used as the
foundation for several other libraries in the JAX ecosystem.In addition, DeepMind has open-sourced anecosystem of libraries around
JAXincludingOptaxfor gradient processing and
optimization,RLaxfor RL algorithms, andchexfor reliable code and testing. (Watch
the NeurIPS 2020 JAX Ecosystem at DeepMind talkhere)Citing JAXTo 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.3.13},
  year = {2018},
}In the above bibtex entry, names are in alphabetical order, the version number
is intended to be that fromjax/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 apaper 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 documentationFor details about the JAX API, see thereference documentation.For getting started as a JAX developer, see thedeveloper documentation.