14. JAX#
Ce cours propose une brève introduction à Google JAX.
JAX est une bibliothèque de calcul scientifique haute performance qui fournit
une interface de type NumPy capable de paralléliser automatiquement sur les CPU et les GPU,
un compilateur just-in-time pour accélérer une large gamme d’opérations numériques, et
De plus en plus, JAX maintient et fournit également des routines de calcul scientifique plus spécialisées, telles que celles que l’on trouve à l’origine dans SciPy.
En plus de ce qui est inclus dans Anaconda, ce cours nécessitera les bibliothèques suivantes :
!pip install jax quantecon
Nous utiliserons les importations suivantes
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import quantecon as qe
14.1. JAX comme remplacement de NumPy#
Examinons les similitudes et les différences entre JAX et NumPy.
14.1.1. Similitudes#
Ci-dessus, nous importons jax.numpy as jnp, qui fournit une interface de type NumPy
pour les opérations sur les tableaux.
L’une des caractéristiques attrayantes de JAX est que, dans la mesure du possible, cette interface se conforme à l’API de NumPy.
En conséquence, nous pouvons souvent utiliser JAX comme un remplacement direct de NumPy.
Voici quelques opérations standard sur les tableaux utilisant jnp :
a = jnp.asarray((1.0, 3.2, -1.5))
print(a)
[ 1. 3.2 -1.5]
print(jnp.sum(a))
2.6999998
print(jnp.dot(a, a))
13.490001
Il faut cependant se rappeler que l’objet tableau a n’est pas un tableau NumPy :
a
Array([ 1. , 3.2, -1.5], dtype=float32)
type(a)
jaxlib._jax.ArrayImpl
Même les applications à valeurs scalaires sur les tableaux renvoient des tableaux JAX plutôt que des scalaires !
jnp.sum(a)
Array(2.6999998, dtype=float32)
14.1.2. Différences#
Examinons maintenant quelques différences entre les opérations sur les tableaux de JAX et de NumPy.
14.1.2.1. La vitesse !#
Une différence majeure est que JAX est plus rapide — et parfois beaucoup plus rapide.
Pour illustrer cela, supposons que nous voulions évaluer la fonction cosinus en de nombreux points.
n = 50_000_000
x = np.linspace(0, 10, n) # Tableau NumPy
14.1.2.1.1. Avec NumPy#
Essayons avec NumPy
with qe.Timer():
# Premier chronométrage NumPy
y = np.cos(x)
0.5685 seconds elapsed
Et encore une fois.
with qe.Timer():
# Deuxième chronométrage NumPy
y = np.cos(x)
0.5706 seconds elapsed
Ici
NumPy utilise un binaire pré-compilé pour appliquer le cosinus à un tableau de flottants
Le binaire s’exécute sur le CPU de la machine locale
14.1.2.1.2. Avec JAX#
Essayons maintenant avec JAX.
x = jnp.linspace(0, 10, n)
Chronométrons la même procédure.
with qe.Timer():
# Première exécution
y = jnp.cos(x)
# Maintient l'interpréteur jusqu'à ce que l'opération sur le tableau soit terminée
y.block_until_ready()
0.1233 seconds elapsed
Note
Ci-dessus, la méthode block_until_ready
maintient l’interpréteur jusqu’à ce que les résultats du calcul soient renvoyés.
Ceci est nécessaire pour chronométrer l’exécution, car JAX utilise une répartition asynchrone, qui
permet à l’interpréteur Python de prendre de l’avance sur les calculs numériques.
Chronométrons-le de nouveau.
with qe.Timer():
# Deuxième exécution
y = jnp.cos(x)
# Maintient l'interpréteur
y.block_until_ready()
0.1029 seconds elapsed
Sur un GPU, ce code s’exécute beaucoup plus rapidement que son équivalent NumPy.
De plus, en général, la deuxième exécution est plus rapide que la première grâce à la compilation JIT.
Cela est dû au fait que même les fonctions intégrées comme jnp.cos sont compilées en JIT — et la
première exécution inclut le temps de compilation.
Pourquoi JAX voudrait-il compiler en JIT des fonctions intégrées comme jnp.cos au lieu de
simplement fournir des versions pré-compilées, comme NumPy ?
La raison est que le compilateur JIT veut se spécialiser sur la taille du tableau utilisé (ainsi que sur le type de données).
La taille importe pour générer du code optimisé, car une parallélisation efficace nécessite d’adapter la taille de la tâche au matériel disponible.
14.1.2.2. Expérience sur la taille#
Nous pouvons vérifier l’affirmation selon laquelle JAX se spécialise sur la taille du tableau en changeant la taille de l’entrée et en observant les temps d’exécution.
x = jnp.linspace(0, 10, n + 1)
with qe.Timer():
# Première exécution
y = jnp.cos(x)
# Maintient l'interpréteur
y.block_until_ready()
0.1404 seconds elapsed
with qe.Timer():
# Deuxième exécution
y = jnp.cos(x)
# Maintient l'interpréteur
y.block_until_ready()
0.0985 seconds elapsed
Le temps d’exécution augmente puis diminue à nouveau (ceci sera plus évident sur le GPU).
Ceci est conforme à la discussion ci-dessus – la première exécution après le changement de taille du tableau montre la surcharge de compilation.
Une discussion plus approfondie de la compilation JIT est fournie ci-dessous.
14.1.2.3. Précision#
Une autre différence entre NumPy et JAX est que JAX utilise par défaut des flottants sur 32 bits.
Cela est dû au fait que JAX est souvent utilisé pour le calcul sur GPU, et la plupart des calculs GPU utilisent des flottants sur 32 bits.
L’utilisation de flottants sur 32 bits peut entraîner des gains de vitesse significatifs avec une faible perte de précision.
Cependant, pour certains calculs, la précision importe.
Dans ces cas, on peut imposer les flottants sur 64 bits via la commande
jax.config.update("jax_enable_x64", True)
Vérifions que cela fonctionne :
jnp.ones(3)
Array([1., 1., 1.], dtype=float64)
14.1.2.4. Immuabilité#
En tant que remplacement de NumPy, une différence plus significative est que les tableaux sont traités comme immuables.
Par exemple, avec NumPy nous pouvons écrire
a = np.linspace(0, 1, 3)
a
array([0. , 0.5, 1. ])
puis muter les données en mémoire :
a[0] = 1
a
array([1. , 0.5, 1. ])
Dans JAX, cela échoue 😱.
a = jnp.linspace(0, 1, 3)
a
Array([0. , 0.5, 1. ], dtype=float64)
try:
a[0] = 1
except Exception as e:
print(e)
JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html
Les concepteurs de JAX ont choisi de rendre les tableaux immuables parce que
JAX utilise un style de programmation fonctionnelle et
la programmation fonctionnelle évite généralement les données mutables
Nous discutons de ces idées ci-dessous.
14.1.2.5. Une solution de contournement#
JAX fournit bien une alternative directe à la modification de tableau sur place
via la méthode at.
a = jnp.linspace(0, 1, 3)
Appliquer at[0].set(1) renvoie une nouvelle copie de a dont le premier élément est fixé à 1
a = a.at[0].set(1)
a
Array([1. , 0.5, 1. ], dtype=float64)
Évidemment, l’utilisation de at présente des inconvénients :
La syntaxe est lourde et
nous voulons éviter de créer de nouveaux tableaux en mémoire à chaque fois que nous changeons une seule valeur !
Par conséquent, la plupart du temps, nous essayons d’éviter cette syntaxe.
(Bien qu’elle puisse en fait être efficace à l’intérieur des fonctions compilées en JIT – mais laissons cela de côté pour l’instant.)
14.2. Programmation fonctionnelle#
D’après la documentation de JAX :
Lorsque vous vous promenez dans la campagne italienne, les gens n’hésiteront pas à vous dire que JAX a « una anima di pura programmazione funzionale ».
En d’autres termes, JAX suppose un style de programmation fonctionnelle.
14.2.1. Fonctions pures#
L’implication majeure est que les fonctions JAX doivent être pures.
Les fonctions pures ont les caractéristiques suivantes :
Déterministes
Sans effets secondaires
Déterministe signifie
Même entrée \(\implies\) même sortie
Les sorties ne dépendent pas de l’état global
En particulier, les fonctions pures renverront toujours le même résultat si elles sont invoquées avec les mêmes entrées.
Sans effets secondaires signifie que la fonction
Ne modifiera pas l’état global
Ne modifiera pas les données passées à la fonction (données immuables)
14.2.2. Exemples – pures et impures#
Voici un exemple d’une fonction impure
tax_rate = 0.1
def add_tax(prices):
for i, price in enumerate(prices):
prices[i] = price * (1 + tax_rate)
prices = [10.0, 20.0]
add_tax(prices)
prices
[11.0, 22.0]
Cette fonction n’est pas pure parce que
effets secondaires — elle modifie la variable globale
pricesnon déterministe — un changement de la variable globale
tax_ratemodifiera les sorties de la fonction, même avec le même tableau d’entréeprices.
Voici une version pure
def add_tax_pure(prices, tax_rate):
new_prices = [price * (1 + tax_rate) for price in prices]
return new_prices
tax_rate = 0.1
prices = (10.0, 20.0)
after_tax_prices = add_tax_pure(prices, tax_rate)
after_tax_prices
[11.0, 22.0]
Elle est pure parce que
toutes les dépendances sont explicites via les arguments de la fonction
et elle ne modifie aucun état externe
14.2.3. Pourquoi la programmation fonctionnelle ?#
Chez QuantEcon, nous adorons les fonctions pures parce qu’elles
Facilitent les tests : chaque fonction peut fonctionner de manière isolée
Favorisent un comportement déterministe et donc la reproductibilité
Empêchent les bugs qui découlent de la mutation d’un état partagé
Le compilateur JAX adore les fonctions pures et la programmation fonctionnelle parce que
Les dépendances de données sont explicites, ce qui aide à optimiser les calculs complexes
Les fonctions pures sont plus faciles à différencier (autodiff)
Les fonctions pures sont plus faciles à paralléliser et à optimiser (elles ne dépendent pas d’un état mutable partagé)
Une autre façon de voir cela est la suivante :
JAX représente les fonctions sous forme de graphes de calcul, qui sont ensuite compilés ou transformés (par ex., différenciés)
Ces graphes de calcul décrivent comment un ensemble donné d’entrées est transformé en une sortie.
Les graphes de calcul de JAX sont purs par construction.
JAX utilise un style de programmation fonctionnelle afin que les fonctions construites par l’utilisateur se traduisent directement dans les représentations théoriques des graphes prises en charge par JAX.
14.3. Nombres aléatoires#
La génération de nombres aléatoires dans JAX diffère significativement des schémas que l’on trouve dans NumPy ou MATLAB.
14.3.1. Approche NumPy / MATLAB#
Dans NumPy / MATLAB, la génération fonctionne en maintenant un état global caché.
np.random.seed(42)
print(np.random.randn(2))
[ 0.49671415 -0.1382643 ]
Chaque fois que nous appelons une fonction aléatoire, l’état caché est mis à jour :
print(np.random.randn(2))
[0.64768854 1.52302986]
Cette fonction n’est pas pure parce que :
Elle est non déterministe : mêmes entrées, sorties différentes
Elle a des effets secondaires : elle modifie l’état du générateur de nombres aléatoires global
Cela est dangereux sous parallélisation — il faut contrôler soigneusement ce qui se passe dans chaque thread.
14.3.2. JAX#
Dans JAX, l’état du générateur de nombres aléatoires est contrôlé explicitement.
Nous produisons d’abord une clé, qui initialise le générateur de nombres aléatoires.
seed = 1234
key = jax.random.key(seed)
Nous pouvons maintenant utiliser la clé pour générer des nombres aléatoires :
x = jax.random.normal(key, (3, 3))
x
Array([[-0.54019824, 0.43957585, -0.01978102],
[ 0.90665474, -0.90831359, 1.32846635],
[ 0.20408174, 0.93096529, 3.30373914]], dtype=float64)
Si nous utilisons à nouveau la même clé, nous initialisons à la même graine, donc les nombres aléatoires sont les mêmes :
jax.random.normal(key, (3, 3))
Array([[-0.54019824, 0.43957585, -0.01978102],
[ 0.90665474, -0.90831359, 1.32846635],
[ 0.20408174, 0.93096529, 3.30373914]], dtype=float64)
Pour produire un tirage (quasi-) indépendant, une option est de « diviser » (split) la clé existante :
key, subkey = jax.random.split(key)
jax.random.normal(key, (3, 3))
Array([[ 1.24104247, 0.12018902, -2.23990047],
[ 0.70507261, -0.85702845, -1.24582014],
[ 0.38454486, 1.32117717, 0.56866901]], dtype=float64)
jax.random.normal(subkey, (3, 3))
Array([[ 0.07627173, -1.30349831, 0.86524323],
[-0.75550773, 0.63958052, 0.47052126],
[-1.72866044, -1.14696564, -1.23328892]], dtype=float64)
Le diagramme suivant illustre comment split produit un arbre de clés à partir d’une
racine unique, chaque clé générant des tirages aléatoires indépendants.
Cette syntaxe semblera inhabituelle pour un utilisateur de NumPy ou Matlab — mais elle prendra plus de sens lorsque nous aborderons la programmation parallèle.
La fonction ci-dessous produit k matrices n x n aléatoires (quasi-) indépendantes en utilisant split.
def gen_random_matrices(
key, # Clé JAX pour les nombres aléatoires
n=2, # Les matrices seront n x n
k=3 # Nombre de matrices à générer
):
matrices = []
for _ in range(k):
key, subkey = jax.random.split(key)
A = jax.random.uniform(subkey, (n, n))
matrices.append(A)
return matrices
seed = 42
key = jax.random.key(seed)
gen_random_matrices(key)
[Array([[0.74211901, 0.54715578],
[0.05988742, 0.32206803]], dtype=float64),
Array([[0.65877976, 0.57087415],
[0.97301903, 0.10138266]], dtype=float64),
Array([[0.68745522, 0.25974132],
[0.06595873, 0.83589118]], dtype=float64)]
Cette fonction est pure
Déterministe : mêmes entrées, même sortie
Sans effets secondaires : aucun état caché n’est modifié
14.3.3. Avantages#
Comme mentionné ci-dessus, cette explicitude est précieuse :
Reproductibilité : facile de reproduire les résultats en réutilisant les clés
Parallélisation : contrôle de ce qui se passe sur des threads séparés
Débogage : l’absence d’état caché rend le code plus facile à tester
Compatibilité JIT : le compilateur peut optimiser les fonctions pures de manière plus agressive
14.4. Compilation JIT#
Le compilateur just-in-time (JIT) de JAX accélère l’exécution en générant un code machine efficace qui varie selon la taille de la tâche et le matériel.
Nous avons vu la puissance du compilateur JIT de JAX combiné à du matériel parallèle
ci-dessus, lorsque nous avons appliqué cos à un grand tableau.
Ici, nous étudions la compilation JIT pour des fonctions plus complexes
14.4.1. Avec NumPy#
Nous essaierons d’abord avec NumPy, en utilisant
def f(x):
y = np.cos(2 * x**2) + np.sqrt(np.abs(x)) + 2 * np.sin(x**4) - x**2
return y
Exécutons avec un grand x
n = 50_000_000
x = np.linspace(0, 10, n)
with qe.Timer():
# Chronométrer le code NumPy
y = f(x)
2.1999 seconds elapsed
Modèle d’exécution impatient (eager)
Chaque opération est exécutée immédiatement dès qu’elle est rencontrée, matérialisant son résultat avant que l’opération suivante ne commence.
Inconvénients
Parallélisation minimale
Empreinte mémoire lourde — produit de nombreux tableaux intermédiaires
Beaucoup de lectures/écritures en mémoire
14.4.2. Avec JAX#
En première approche, nous remplaçons np par jnp partout :
def f(x):
y = jnp.cos(2 * x**2) + jnp.sqrt(jnp.abs(x)) + 2 * jnp.sin(x**4) - x**2
return y
x = jnp.linspace(0, 10, n)
Chronométrons-le maintenant.
with qe.Timer():
# Premier appel
y = f(x)
# Maintient l'interpréteur
jax.block_until_ready(y);
0.9578 seconds elapsed
with qe.Timer():
# Deuxième appel
y = f(x)
# Maintient l'interpréteur
jax.block_until_ready(y);
0.8537 seconds elapsed
Le résultat est similaire à l’exemple cos — JAX est plus rapide, en particulier à la
deuxième exécution après la compilation JIT.
Cela est dû au fait que les opérations individuelles sur les tableaux sont parallélisées sur le GPU
Mais nous utilisons toujours l’exécution impatiente
beaucoup de mémoire à cause des tableaux intermédiaires
beaucoup de lectures/écritures en mémoire
De plus, de nombreux noyaux séparés sont lancés sur le GPU
14.4.3. Compiler la fonction entière#
Heureusement, avec JAX, nous avons un autre tour dans notre sac — nous pouvons compiler en JIT la fonction entière, et pas seulement les opérations individuelles.
Le compilateur fusionne toutes les opérations sur les tableaux en un seul noyau optimisé
Essayons ceci avec la fonction f :
f_jax = jax.jit(f)
with qe.Timer():
# Première exécution
y = f_jax(x)
# Maintient l'interpréteur
jax.block_until_ready(y);
0.5731 seconds elapsed
with qe.Timer():
# Deuxième exécution
y = f_jax(x)
# Maintient l'interpréteur
jax.block_until_ready(y);
0.5426 seconds elapsed
Le temps d’exécution s’est encore amélioré — maintenant parce que nous avons fusionné toutes les opérations
Optimisation agressive basée sur la séquence de calcul entière
Élimination des multiples appels à l’accélérateur matériel
L’empreinte mémoire est également bien plus faible — aucune création de tableaux intermédiaires
Au passage, une syntaxe plus courante lorsqu’on cible une fonction pour le compilateur JIT est
@jax.jit
def f(x):
pass # placer le corps de la fonction ici
14.4.5. Compilation de fonctions non pures#
Bien que JAX ne génère généralement pas d’erreurs lors de la compilation de fonctions impures, l’exécution devient imprévisible !
Voici une illustration de ce fait :
a = 1 # global
@jax.jit
def f(x):
return a + x
x = jnp.ones(2)
f(x)
Array([2., 2.], dtype=float64)
Dans le code ci-dessus, la valeur globale a=1 est fusionnée dans la fonction jittée.
Même si nous changeons a, la sortie de f ne sera pas affectée — tant que la même version compilée est appelée.
a = 42
f(x)
Array([2., 2.], dtype=float64)
Changer la dimension de l’entrée déclenche une nouvelle compilation de la fonction, moment auquel le changement de la valeur de a prend effet :
x = jnp.ones(3)
f(x)
Array([43., 43., 43.], dtype=float64)
Morale de l’histoire : écrivez des fonctions pures lorsque vous utilisez JAX !
14.5. Vectorisation avec vmap#
Une autre transformation puissante de JAX est jax.vmap, qui vectorise automatiquement
une fonction écrite pour une entrée unique afin qu’elle opère sur des
lots.
Cela évite d’avoir à écrire manuellement du code vectorisé ou d’utiliser des boucles explicites.
14.5.1. Un exemple simple#
Supposons que nous ayons une fonction qui calcule la différence entre la moyenne et la médiane pour un tableau de nombres.
def mm_diff(x):
return jnp.mean(x) - jnp.median(x)
Nous pouvons l’appliquer à un seul vecteur :
x = jnp.array([1.0, 2.0, 5.0])
mm_diff(x)
Array(0.66666667, dtype=float64)
Supposons maintenant que nous ayons une matrice et que nous voulions calculer ces statistiques pour chaque ligne.
Sans vmap, nous aurions besoin d’une boucle explicite :
X = jnp.array([[1.0, 2.0, 5.0],
[4.0, 5.0, 6.0],
[1.0, 8.0, 9.0]])
for row in X:
print(mm_diff(row))
0.6666666666666665
0.0
-2.0
Cependant, les boucles Python sont lentes et ne peuvent pas être compilées ou parallélisées efficacement par JAX.
Avec vmap, nous pouvons éviter les boucles et garder le calcul sur l’accélérateur :
batch_mm_diff = jax.vmap(mm_diff) # Crée une nouvelle version « vectorisée »
batch_mm_diff(X) # Applique à chaque ligne de X
Array([ 0.66666667, 0. , -2. ], dtype=float64)
14.5.2. Combinaison des transformations#
L’une des forces de JAX est que les transformations se composent naturellement.
Par exemple, nous pouvons compiler en JIT une fonction vectorisée :
fast_batch_mm_diff = jax.jit(jax.vmap(mm_diff))
fast_batch_mm_diff(X)
Array([ 6.66666667e-01, -2.77555756e-16, -2.00000000e+00], dtype=float64)
Cette composition de jit, vmap, et (comme nous le verrons ensuite) grad est centrale dans
la conception de JAX et le rend particulièrement puissant pour le calcul scientifique et
l’apprentissage automatique.
14.6. Différenciation automatique : un aperçu#
JAX peut utiliser la différenciation automatique pour calculer des gradients.
Cela peut être extrêmement utile pour l’optimisation et la résolution de systèmes non linéaires.
Voici une illustration simple faisant intervenir la fonction \(f(x) = x^2 / 2\) :
def f(x):
return (x**2) / 2
f_prime = jax.grad(f)
f_prime(10.0)
Array(10., dtype=float64, weak_type=True)
Traçons la fonction et sa dérivée, en notant que \(f'(x) = x\).
fig, ax = plt.subplots()
x_grid = jnp.linspace(-4, 4, 200)
ax.plot(x_grid, f(x_grid), label="$f$")
ax.plot(x_grid, [f_prime(x) for x in x_grid], label="$f'$")
ax.legend(loc='upper center')
plt.show()
La différenciation automatique est un sujet vaste avec de nombreuses applications en économie et en finance. Nous en fournissons un traitement plus approfondi dans notre cours sur l’autodiff.
14.7. Exercices#
Exercice 14.1
Dans la section Exercices de notre cours sur Numba, nous avons utilisé Monte-Carlo pour évaluer une option d’achat européenne.
Le code a été accéléré par le multithreading basé sur Numba.
Essayez d’écrire une version de cette opération pour JAX, en utilisant tous les mêmes paramètres.
Solution
Voici une solution :
M = 10_000_000
n, β, K = 20, 0.99, 100
μ, ρ, ν, S0, h0 = 0.0001, 0.1, 0.001, 10, 0
@jax.jit
def compute_call_price_jax(β=β,
μ=μ,
S0=S0,
h0=h0,
K=K,
n=n,
ρ=ρ,
ν=ν,
M=M,
key=jax.random.key(1)):
s = jnp.full(M, np.log(S0))
h = jnp.full(M, h0)
def update(i, loop_state):
s, h, key = loop_state
key, subkey = jax.random.split(key)
Z = jax.random.normal(subkey, (2, M))
s = s + μ + jnp.exp(h) * Z[0, :]
h = ρ * h + ν * Z[1, :]
new_loop_state = s, h, key
return new_loop_state
initial_loop_state = s, h, key
final_loop_state = jax.lax.fori_loop(0, n, update, initial_loop_state)
s, h, key = final_loop_state
expectation = jnp.mean(jnp.maximum(jnp.exp(s) - K, 0))
return β**n * expectation
Note
Nous utilisons jax.lax.fori_loop au lieu d’une boucle for Python.
Cela permet à JAX de compiler la boucle efficacement sans la dérouler,
ce qui réduit considérablement le temps de compilation pour les grands tableaux.
Exécutons-le une fois pour le compiler :
with qe.Timer():
compute_call_price_jax().block_until_ready()
5.6938 seconds elapsed
Et maintenant chronométrons-le :
with qe.Timer():
compute_call_price_jax().block_until_ready()
5.4274 seconds elapsed
14.4.4. Comment fonctionne la compilation JIT#
Lorsque nous appliquons
jax.jità une fonction, JAX la trace : au lieu d’exécuter les opérations immédiatement, il enregistre la séquence d’opérations sous forme de graphe de calcul et transmet ce graphe au compilateur XLA.XLA fusionne et optimise ensuite les opérations en un seul noyau compilé adapté au matériel disponible (CPU, GPU ou TPU).
Le premier appel à une fonction compilée en JIT entraîne une surcharge de compilation, mais les appels ultérieurs avec les mêmes formes et types d’entrée réutilisent le code compilé mis en cache et s’exécutent à pleine vitesse.