jax-0.4.14.tar.gz

<div align="center"> <img src="https://raw.githubusercontent.com/google/jax/main/images/jax_logo_250px.png" alt="logo"></img> </div> # JAX: Autograd and XLA   [**Quickstart**](#quickstart-colab-in-the-cloud) | [**Transformations**](#transformations) | [**Install guide**](#installation) | [**Neural net libraries**](#neural-network-libraries) | [**Change logs**](https://jax.readthedocs.io/en/latest/changelog.html) | [**Reference docs**](https://jax.readthedocs.io/en/latest/) ## What is JAX? JAX is [Autograd](https://github.com/hips/autograd) and [XLA](https://www.tensorflow.org/xla), brought together for high-performance machine learning research. With its updated version of [Autograd](https://github.com/hips/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`](#automatic-differentiation-with-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](https://www.tensorflow.org/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-with-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`](#spmd-programming-with-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](#transformations). Both [`grad`](#automatic-differentiation-with-grad) and [`jit`](#compilation-with-jit) are instances of such transformations. Others are [`vmap`](#auto-vectorization-with-vmap) for automatic vectorization and [`pmap`](#spmd-programming-with-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](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html). Please help by trying it out, [reporting bugs](https://github.com/google/jax/issues), and letting us know what you think! ```python 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) # 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 = jit(grad(loss)) # compiled gradient evaluation function perex_grads = jit(vmap(grad_loss, in_axes=(None, 0, 0))) # fast per-example grads ``` ### Contents * [Quickstart: Colab in the Cloud](#quickstart-colab-in-the-cloud) * [Transformations](#transformations) * [Current gotchas](#current-gotchas) * [Installation](#installation) * [Neural net libraries](#neural-network-libraries) * [Citing JAX](#citing-jax) * [Reference documentation](#reference-documentation) ## 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](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html) - [Training a Simple Neural Network, with TensorFlow Dataset Data Loading](https://colab.research.google.com/github/google/jax/blob/main/docs/notebooks/neural_network_with_tfds_data.ipynb) **JAX now runs on Cloud TPUs.** To try out the preview, see the [Cloud TPU Colabs](https://github.com/google/jax/tree/main/cloud_tpu_colabs). For a deeper dive into JAX: - [The Autodiff Cookbook, Part 1: easy and powerful automatic differentiation in JAX](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html) - [Common gotchas and sharp edges](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html) - See the [full list of notebooks](https://github.com/google/jax/tree/main/docs/notebooks). You can also take a look at [the mini-libraries in `jax.example_libraries`](https://github.com/google/jax/tree/main/jax/example_libraries/README.md), like [`stax` for building neural networks](https://github.com/google/jax/tree/main/jax/example_libraries/README.md#neural-net-building-with-stax) and [`optimizers` for first-order stochastic optimization](https://github.com/google/jax/tree/main/jax/example_libraries/README.md#first-order-optimization), or the [examples](https://github.com/google/jax/tree/main/examples). ## Transformations At its core, JAX is an extensible system for transforming numerical functions. Here are four transformations of primary interest: `grad`, `jit`, `vmap`, and `pmap`. ### Automatic differentiation with `grad` JAX has roughly the same API as [Autograd](https://github.com/hips/autograd). The most popular function is [`grad`](https://jax.readthedocs.io/en/latest/jax.html#jax.grad) for reverse-mode gradients: ```python 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`. ```python print(grad(grad(grad(tanh)))(1.0)) # prints 0.62162673 ``` For more advanced autodiff, you can use [`jax.vjp`](https://jax.readthedocs.io/en/latest/jax.html#jax.vjp) for reverse-mode vector-Jacobian products and [`jax.jvp`](https://jax.readthedocs.io/en/latest/jax.html#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](https://jax.readthedocs.io/en/latest/_autosummary/jax.hessian.html#jax.hessian): ```python from jax import jit, jacfwd, jacrev def hessian(fun): return jit(jacfwd(jacrev(fun))) ``` As with [Autograd](https://github.com/hips/autograd), you're free to use differentiation with Python control structures: ```python 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](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation) and the [JAX Autodiff Cookbook](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html) for more. ### Compilation with `jit` You can use XLA to compile your functions end-to-end with [`jit`](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit), used either as an `@jit` decorator or as a higher-order function. ```python 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