G1Joshi/Agent-Skills

JAX

JAX, c'est « NumPy sous stéroïdes ». Il combine Autograd (différentiation automatique) avec XLA (compilation). En 2025, Flax NNX (programmation orientée objet à la PyTorch) devient la norme.

Quand l'utiliser

• Calcul numérique et scientifique haute performance : Algèbre linéaire accélérée sur GPU/TPU grâce à Autograd et XLA.
• Recherche personnalisée en deep learning : Construction de réseaux de neurones avec Flax, Haiku ou Equinox, sans lourdeur de framework.
• Transformations de fonctions composables : Enchaînement transparent de jit (compilation), grad (dérivées) et vmap (vectorisation).
• Entraînement distribué multi-périphériques : Parallélisation des calculs sur des clusters avec pmap et les shard maps (jax.experimental.shard_map).

Démarrage rapide

import jax
import jax.numpy as jnp

# Multiplication matricielle compilée JIT et différentiation automatique
@jax.jit
def loss_fn(w, x, y):
    pred = jnp.dot(x, w)
    return jnp.mean((pred - y) ** 2)

# Calcul simultané de la valeur et des gradients
grad_fn = jax.grad(loss_fn)

w = jnp.array([1.0, 2.0])
x = jnp.array([[1.0, 0.5], [2.0, 1.0]])
y = jnp.array([2.0, 4.0])

grads = grad_fn(w, x, y)
print("Gradients :", grads)

Concepts fondamentaux

#Transformations de fonctions composables : jit, grad et vmap

Combinaison de la compilation juste-à-temps, de la différentiation automatique et du traitement par lots automatisé :

import jax
import jax.numpy as jnp

# Définition d'une fonction pure
def loss_fn(w, x, y):
    pred = jnp.dot(x, w)
    return jnp.mean((pred - y) ** 2)

# Transformation 1 : Gradient automatique par rapport aux poids
grad_fn = jax.grad(loss_fn)

# Transformation 2 : Compilation JIT via XLA pour une accélération GPU extrême
fast_grad_fn = jax.jit(grad_fn)

# Transformation 3 : Vectorisation automatique sur les lots
# batch_loss_fn gère les matrices 2D automatiquement sans boucles
vmapped_loss = jax.vmap(loss_fn, in_axes=(None, 0, 0))

# Exemple de calcul
key = jax.random.PRNGKey(42)
w = jax.random.normal(key, (3,))
x = jnp.array([[1.0, 2.0, 3.0]])
y = jnp.array([5.0])

grad_val = fast_grad_fn(w, x, y)
print("Gradient calculé :", grad_val)

#Gestion explicite des clés PRNG

Génération de nombres pseudo-aléatoires sans état :

import jax

# Initialisation de la clé PRNG racine
key = jax.random.PRNGKey(2026)

# Les clés JAX doivent être divisées explicitement ; jamais réutilisées
key, subkey1, subkey2 = jax.random.split(key, 3)

data_normal = jax.random.normal(subkey1, shape=(4, 4))
data_uniform = jax.random.uniform(subkey2, shape=(4, 4))

print("Moyenne de l'échantillon normal :", jnp.mean(data_normal))

#Modèle avec état utilisant Equinox et des fonctions pures

Réseaux de neurones expressifs et orientés objet construits sur des fonctions JAX pures :

import equinox as eqx
import jax
import jax.numpy as jnp

class SimpleMLP(eqx.Module):
    linear1: eqx.nn.Linear
    linear2: eqx.nn.Linear

    def __init__(self, in_features, out_features, key):
        k1, k2 = jax.random.split(key)
        self.linear1 = eqx.nn.Linear(in_features, 64, key=k1)
        self.linear2 = eqx.nn.Linear(64, out_features, key=k2)

    def __call__(self, x):
        x = jax.nn.relu(self.linear1(x))
        return self.linear2(x)

key = jax.random.PRNGKey(0)
model = SimpleMLP(in_features=10, out_features=2, key=key)
sample_input = jnp.ones((10,))
output = model(sample_input)
print("Sortie du modèle :", output)

Schémas courants

Vectorisation par lots avec vmap

Problème : L'écriture de boucles Python manuelles pour évaluer des fonctions sur des lots ralentit l'exécution.

Solution : Utilisez jax.vmap pour vectoriser automatiquement les fonctions mono-exemple :

def predict_single(w, x):
    return jnp.dot(x, w)

# Vectorisation automatique sur la dimension du lot (axe 0 de x)
predict_batch = jax.vmap(predict_single, in_axes=(None, 0))

batch_x = jnp.ones((100, 10))
weights = jnp.ones(10)
predictions = predict_batch(weights, batch_x) # Forme de sortie : (100,)

Bonnes pratiques (2026)

• À faire : Écrire du code purement fonctionnel : aucun effet de bord, aucune mutation de tableau en place (x.at[idx].set(val) plutôt que x[idx] = val).
• À faire : Encapsuler les fonctions critiques en performance avec @jax.jit pour déclencher la compilation XLA.
• À faire : Diviser explicitement les clés PRNG (jax.random.split(key)) à chaque génération de nombres aléatoires.
• À faire : Utiliser des bibliothèques modernes de haut niveau comme Equinox ou Flax Linen plutôt que d'écrire des dictionnaires de paramètres bruts.
• À ne pas faire : Utiliser des conditionnelles Python standard (if x > 0:) à l'intérieur de fonctions JIT sur des traceurs dynamiques ; utiliser jax.lax.cond.
• À ne pas faire : Réutiliser les clés PRNG ; la réutilisation génère des nombres aléatoires statistiquement corrélés.
• À ne pas faire : Modifier des variables globales à l'intérieur de fonctions transformées avec jit, grad ou vmap.

Dépannage

Erreur	Cause	Solution
ConcretizationTypeError: Abstract tracer value used in Python if/while	Une branche conditionnelle Python dépend de valeurs JAX dynamiques dans jit.	Utiliser jax.lax.cond ou jax.lax.while_loop au lieu du if Python standard.
JAX arrays are immutable	Tentative d'affectation en place (arr[0] = 5).	Utiliser la syntaxe de mise à jour fonctionnelle : arr = arr.at[0].set(5).
PRNGKey reuse warning / identical randomness	La réutilisation d'une même clé PRNG produit des nombres aléatoires identiques.	Diviser les clés avant chaque opération aléatoire : key, subkey = jax.random.split(key).

Références

• Documentation JAX