How to Turn Familiar NumPy into Compilable GPU Code

If you've ever tried to speed up heavy computations in Python, you've likely hit the limitations of the standard stack. NumPy is fast thanks to its C backend, but can't work with GPUs out of the box. PyTorch and TensorFlow solve this problem, but bring along bulky abstractions, layer classes, and their own computation graph semantics.

In 2018, engineers at Google open-sourced the JAX library. The idea under the hood is simple: give developers familiar NumPy syntax, but add automatic differentiation and the XLA compiler on top. The result is clean functional Python that gets compiled on-the-fly into optimized machine code for GPUs or TPUs.

What JAX Actually Is

Many people think of JAX as just another ML framework, but that's not quite accurate. The authors themselves state in the repository that it's a system of composable functional transformations for arrays.

Instead of building complex object models, JAX encourages working with pure functions. You write regular Python code, and then apply transformer functions to it. No hidden global states or in-place data mutations. If you need to compute a gradient, compile a piece of the program, or distribute computations across batches, you just wrap the function in the appropriate decorator.

Let's look at the four main transformations that the entire library is built around.

Automatic Differentiation via grad

The jax.grad function takes your function and returns a new one that computes the gradient of the original:

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))  # 0.4199743

You can compute gradients of any order by simply nesting jax.grad calls. The algorithm handles standard Python conditionals if/else, loops, and recursion without issues.

Compilation via jit

Regular Python executes each array operation sequentially, with overhead from function calls and intermediate memory allocation. The jax.jit decorator sends the function body to the XLA compiler:

def calculate(x):
    return x * x + x * 2.0

x = jnp.ones((5000, 5000))

# Обычный запуск против скомпилированного
fast_calculate = jax.jit(calculate)

The compiler fuses elementary operations into a single compute kernel. As a result, data doesn't bounce back and forth between cache and GPU memory, and performance increases by orders of magnitude.

Code Vectorization via vmap

Everyone who's written machine learning algorithms has spent hours adjusting tensor dimensions to match the batch size. jax.vmap solves this headache: you write logic for a single element or vector, and JAX itself vectorizes the operation:

def l1_distance(x, y):
    return jnp.sum(jnp.abs(x - y))

# Превращаем функцию для векторов в функцию для матриц
def pairwise_distances(xs):
    return jax.vmap(jax.vmap(l1_distance, (0, None)), (None, 0))(xs, xs)

xs = jax.random.normal(jax.random.key(0), (100, 3))
matrix = pairwise_distances(xs)  # форма (100, 100)

Instead of a slow Python loop, the vectorizer pushes the loop inside low-level operations, turning matrix-vector multiplications into full matrix multiplications.

Parallelism and Data Sharding

When a model no longer fits in a single accelerator's memory, JAX offers a declarative approach to parallelism. You define a device mesh and array partitioning rules (partition spec), and the compiler itself distributes computations and configures data exchange between cards:

from jax.sharding import set_mesh, AxisType, PartitionSpec as P

# Создаем сетку из 8 ускорителей
mesh = jax.make_mesh((8,), ('data',), axis_types=(AxisType.Explicit,))
set_mesh(mesh)

# Шардируем входные данные
inputs, targets = jax.device_put((inputs, targets), P('data'))

# Обычная функция градиента теперь выполняется параллельно
grad_fn = jax.jit(jax.grad(loss_fn))
grads = grad_fn(params, (inputs, targets))

Caveats and Specifics

JAX has a flip side that you have to get used to.

First, the functional paradigm requires no side effects. You can't just change an array element by index (arr[0] = 5), because arrays in JAX are immutable. For this, use the arr.at[0].set(5) method.

Second, random number generation requires explicit passing of state keys (jax.random.key), because a global seed would break reproducibility during parallel compilation.

Third, debugging JIT code can be unfamiliar: on the first call, the function goes through a tracing phase, and regular Python print statements inside it will only execute once.

Installation and Platforms

The library officially supports Linux and macOS, and also runs on Windows via the WSL2 subsystem.

For running on a regular CPU:

pip install -U jax

For building with NVIDIA CUDA support:

pip install -U "jax[cuda13]"

There's also support for Google TPU accelerators and AMD ROCm.

Is It Worth Trying

JAX is great for research tasks, physics modeling, non-standard optimizations, and scientific computing where PyTorch feels too bulky and plain NumPy isn't fast enough. If your project is hitting performance limits on mathematical operations or requires computing complex derivatives, taking a look at jax-ml/jax is definitely worth it.