Download
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.
Features
- JAX is Autograd and XLA, brought together for high-performance machine learning research
- You can even program multiple GPUs or TPU cores at once using pmap, and differentiate through the whole thing
- JAX is really an extensible system for composable function transformations
- Jump right in using a notebook in your browser, connected to a Google Cloud GPU
- At its core, JAX is an extensible system for transforming numerical functions
- You can use XLA to compile your functions end-to-end
Project Samples
