Hoe je vertrouwde NumPy omzet in compileerbare GPU-code

Als je ooit hebt geprobeerd zware berekeningen in Python te versnellen, ben je waarschijnlijk tegen de beperkingen van de standaard stack aangelopen. NumPy is snel dankzij de C-backend, maar werkt niet standaard met GPU's. PyTorch en TensorFlow lossen dit probleem op, maar brengen logge abstracties, laagklassen en hun eigen berekeningsgraaf-semantiek met zich mee.

In 2018 hebben engineers bij Google de JAX-bibliotheek open-source gemaakt. Het idee erachter is eenvoudig: geef ontwikkelaars vertrouwde NumPy-syntax, maar voeg automatische diffusentiatie en de XLA-compiler toe. Het resultaat is schone functionele Python die on-the-fly wordt gecompileerd naar geoptimaliseerde machinecode voor GPU's of TPU's.

Wat JAX eigenlijk is

Veel mensen beschouwen JAX als slechts een andere ML-framework, maar dat is niet helemaal accuraat. De auteurs zelf stellen in de repository dat het een systeem is van samenstelbare functionele transformaties voor arrays.

In plaats van complexe objectmodellen te bouwen, moedigt JAX het werken met pure functies aan. Je schrijft reguliere Python-code en past vervolgens transformatorfuncties toe. Geen verborgen globale staten of in-place datamutaties. Als je een gradiënt moet berekenen, een deel van het programma moet compileren, of berekeningen over batches moet verdelen, wikkel je de functie gewoon in de juiste decorator.

Laten we naar de vier hoofdtransformaties kijken waarop de hele bibliotheek is gebouwd.

Automatische diffusentiatie via grad

De jax.grad functie neemt je functie en retourneert een nieuwe die de gradiënt van de originele berekent:

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

Je kunt gradiënten van elke orde berekenen door eenvoudig jax.grad aanroepen te nesten. Het algoritme verwerkt standaard Python-conditionals if/else, loops en recursie zonder problemen.

Compilatie via jit

Reguliere Python voert elke array-bewerking sequentieel uit, met overhead van functie-aanroepen en tussentijdse geheugentoewijzing. De jax.jit decorator stuurt de functiebody naar de XLA-compiler:

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

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

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

De compiler fuseert elementaire bewerkingen tot één compute kernel. Als gevolg daarvan stuitert data niet heen en weer tussen cache en GPU-geheugen, en neemt de prestatie met ordes van grootte toe.

Code-vectorisatie via vmap

Iedereen die machine learning-algoritmen heeft geschreven, heeft uren besteed aan het aanpassen van tensor-dimensies om overeen te komen met de batchgrootte. jax.vmap lost deze hoofdpijn op: je schrijft logica voor een enkel element of vector, en JAX vectoriseert de bewerking zelf:

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)

In plaats van een langzame Python-loop, duwt de vectoriseerder de loop in low-level bewerkingen, en verandert matrix-vector vermenigvuldigingen in volledige matrixvermenigvuldigingen.

Parallellisme en data sharding

Wanneer een model niet meer in het geheugen van één accelerator past, biedt JAX een declaratieve aanpak voor parallellisme. Je definieert een device mesh en array-partitieregels (partition spec), en de compiler zelf verdeelt berekeningen en configureert data-uitwisseling tussen kaarten:

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

Valkuilen en specificaties

JAX heeft een keerzijde waar je aan moet wennen.

Ten eerste vereist het functionele paradigma geen neveneffecten. Je kunt niet zomaar een array-element wijzigen via index (arr[0] = 5), omdat arrays in JAX onveranderlijk zijn. Gebruik hiervoor de arr.at[0].set(5) methode.

Ten tweede vereist willekeurige nummergeneratie expliciet doorgeven van status-sleutels (jax.random.key), omdat een globale seed de reproduceerbaarheid tijdens parallelle compilatie zou breken.

Ten derde kan debuggen van JIT-code onwennig zijn: bij de eerste aanroep doorloopt de functie een tracing-fase, en reguliere Python print-statements binnenin worden slechts één keer uitgevoerd.

Installatie en platforms

De bibliotheek ondersteunt officieel Linux en macOS, en draait ook op Windows via het WSL2-subsysteem.

Voor uitvoering op een reguliere CPU:

pip install -U jax

Voor bouwen met NVIDIA CUDA-ondersteuning:

pip install -U "jax[cuda13]"

Er is ook ondersteuning voor Google TPU-accelerators en AMD ROCm.

Is het de moeite waard om te proberen

JAX is geweldig voor onderzoekstaken, fysicamodellering, niet-standaard optimalisaties en wetenschappelijke berekeningen waar PyTorch te log aanvoelt en gewoon NumPy niet snel genoeg is. Als je project prestatielimieten raakt bij wiskundige bewerkingen of complexe afgeleiden moet berekenen, is het zeker de moeite waard om naar jax-ml/jax te kijken.