JAX: The New Kid in Machine Learning Town

In the ever-evolving landscape of machine learning (ML), a new player has emerged that promises to revolutionize the way we approach ML programming: JAX. Developed by Google, JAX is designed to make ML programming more intuitive, structured, and clean. While it may seem like a contender to replace established libraries like TensorFlow and PyTorch, JAX operates on a fundamentally different architecture that sets it apart. As a friend of mine aptly put it, "We had all sorts of Aces, Kings, and Queens. Now we have JAX."

In this article, we will delve into what JAX is, why it stands out among other libraries, and explore its powerful features through practical code snippets. If you’re ready to discover the potential of JAX, let’s dive in!

What is JAX?

JAX is a Python library designed for high-performance ML research. At its core, JAX is a numerical computing library, much like NumPy, but with significant enhancements. It was developed by Google and is used internally by both Google and DeepMind teams. JAX allows users to leverage the power of GPUs and TPUs seamlessly, making it an attractive option for researchers and developers alike.

Source: JAX documentation

Installing JAX

Before we explore the advantages of JAX, it’s essential to install it in your Python environment or Google Colab. You can easily install JAX using pip:

$ pip install --upgrade jax jaxlib

This command will install JAX for CPU execution. If you want to enable GPU support, ensure you have CUDA and cuDNN installed, then run:

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

For troubleshooting, refer to the official GitHub instructions.

Now, let’s import JAX alongside NumPy for comparison:

import jax
import jax.numpy as jnp
import numpy as np

JAX Basics

JAX’s primary purpose is to perform numeric operations in an expressible and high-performance manner. The syntax is almost identical to NumPy, making it easy for users familiar with NumPy to transition to JAX. For example, to create an array of zeros:

x = np.zeros(10)
y = jnp.zeros(10)

The difference lies beneath the surface.

The DeviceArray

One of JAX’s main advantages is its ability to run the same program on hardware accelerators like GPUs and TPUs without any changes. This is achieved through an underlying structure called DeviceArray, which replaces NumPy’s standard array.

DeviceArrays are lazy, meaning they keep values in the accelerator and pull them only when needed. You can use DeviceArrays just like standard arrays, passing them to other libraries, plotting graphs, and performing differentiation seamlessly. Furthermore, JAX supports the majority of NumPy’s API, so your JAX code will closely resemble your existing NumPy code.

Speed Comparison

JAX is designed for speed. To illustrate this, let’s create two arrays of size (1000, 1000) using NumPy and JAX and calculate the inner product with itself. We’ll use the %timeit magic command to measure execution time:

x = np.random.rand(1000, 1000)
y = jnp.array(x)

%timeit -n 1 -r 1 np.dot(x, x)
%timeit -n 1 -r 1 jnp.dot(y, y).block_until_ready()

The results will likely show that JAX outperforms NumPy, especially when utilizing GPU acceleration. Note the use of block_until_ready(), which ensures we wait for the asynchronous execution to complete before measuring time.

Why JAX?

If speed and automatic GPU support aren’t enough to convince you, let’s explore some of JAX’s unique features that set it apart from other libraries.

Automatic Differentiation with grad()

One of JAX’s standout features is its ability to perform automatic differentiation through the grad() function. This capability is invaluable for deep learning applications, as it simplifies backpropagation.

Here’s an example of how to use grad() to differentiate a simple quadratic function:

from jax import grad

def f(x):
    return 3 * x**2 + 2 * x + 5

def f_prime(x):
    return 6 * x + 2

grad(f)(1.0), f_prime(1.0)

JAX computes the derivative analytically under the hood, applying the chain rule to derive the gradient efficiently.

Accelerated Linear Algebra (XLA) Compiler

JAX’s speed is further enhanced by the Accelerated Linear Algebra (XLA) compiler, which optimizes matrix operations by compiling code into a set of computation kernels. This optimization allows JAX to execute operations as quickly as possible.

Just-In-Time Compilation (jit)

To fully leverage XLA, JAX employs Just-In-Time (JIT) compilation. This technique compiles code during execution rather than beforehand, resulting in faster execution times. You can use the jit() function or the @jit decorator to apply JIT compilation:

from jax import jit

x = np.random.rand(1000, 1000)
y = jnp.array(x)

def f(x):
    for _ in range(10):
        x = 0.5 * x + 0.1 * jnp.sin(x)
    return x

g = jit(f)

%timeit -n 5 -r 5 f(y).block_until_ready()
%timeit -n 5 -r 5 g(y).block_until_ready()

The performance improvement is often significant, especially for deep learning tasks where backpropagation can be accelerated by combining jit with grad().

Parallel Computation with pmap

JAX also supports parallel computation across multiple devices using the pmap function. This allows you to distribute computations across available devices automatically:

from jax import pmap

def f(x):
    return jnp.sin(x) + x**2

f(np.arange(4)), pmap(f)(np.arange(4))

When using pmap, the DeviceArray becomes a ShardedDeviceArray, which manages parallel execution across devices.

Automatic Vectorization with vmap

The vmap function enables automatic vectorization, allowing you to apply functions to batches of data efficiently. Here’s an example:

from jax import vmap

def f(x):
    return jnp.square(x)

f(jnp.arange(10)), vmap(f)(jnp.arange(10))

Using vmap, JAX processes the entire batch in a single operation, improving both speed and memory efficiency.

Pseudo-Random Number Generator

JAX’s random number generator operates differently from NumPy’s. Instead of using a standard stateful PseudoRandom Number Generator (PRNG), JAX requires an explicit PRNG state to be passed as the first argument:

from jax import random

key = random.PRNGKey(5)
random.uniform(key)

This design allows for better vectorization and parallel computation.

Asynchronous Dispatch

JAX employs asynchronous dispatch, meaning it doesn’t wait for operations to complete before returning control to the Python program. Instead, it returns a DeviceArray, which acts as a future value that will be produced on an accelerator device. This approach allows Python code to continue executing while the accelerator processes computations.

Profiling JAX

JAX supports profiling through TensorBoard, enabling you to visualize performance metrics. You can also use JAX’s built-in Device Memory Profiler to gain insights into how your code executes on GPUs and TPUs:

import jax
import jax.numpy as jnp
import jax.profiler

def func1(x):
    return jnp.tile(x, 10) * 0.5

def func2(x):
    y = func1(x)
    return y, jnp.tile(x, 10) + 1

x = jax.random.normal(jax.random.PRNGKey(42), (1000, 1000))
y, z = func2(x)
z.block_until_ready()
jax.profiler.save_device_memory_profile("memory.prof")

You can then analyze the profiling data using tools like pprof.

Conclusion

In this article, we explored the powerful features of JAX and how it stands out in the realm of machine learning libraries. From automatic differentiation to accelerated linear algebra and parallel computation, JAX offers a wealth of tools that can enhance your ML workflows.

If you’re interested in experimenting with JAX, you can find the full code in this Colab notebook or in our GitHub repository.

Stay tuned for future articles where we will delve deeper into building and training deep neural networks with JAX, as well as exploring various frameworks built on top of it. If you found this article insightful, please share it on social media!

References

Deep Learning in Production Book 📖

Learn how to build, train, deploy, scale, and maintain deep learning models. Understand ML infrastructure and MLOps using hands-on examples. Learn more.

Disclosure: Please note that some of the links above might be affiliate links, and at no additional cost to you, we will earn a commission if you decide to make a purchase after clicking through.