Introducere în JAX

PyTorch modifică tensori. TensorFlow construiește grafuri. JAX compilează funcții pure. Ultima abordare schimbă modul în care gândiți despre învățarea profundă.

Tip: Construcție Limbaje: Python Cerințe preliminare: Faza 03, lecțiile 01–10; noțiuni de bază NumPy Timp: ~90 de minute

Obiective de învățare

• Scrieți cod de rețele neuronale bazat pe funcții pure, folosind API-ul funcțional JAX (jax.numpy, jax.grad, jax.jit, jax.vmap)
• Explicați diferența esențială de proiectare dintre mutația eager din PyTorch și modelul de compilare funcțională JAX
• Aplicați compilarea JIT și vectorizarea vmap pentru a accelera buclele de antrenare față de codul Python naiv
• Antrenați o rețea simplă în JAX și comparați gestionarea explicită a stării cu abordarea orientată pe obiecte din PyTorch

Problema

Știți să construiți rețele neuronale în PyTorch. Definiți un nn.Module, apelați .backward(), executați pasul optimizatorului. Funcționează. Milioane de oameni îl folosesc.

Dar PyTorch are o constrângere înscrisă în ADN-ul său: urmărește operațiile eager, pe rând, în Python. Fiecare tensor + tensor este o lansare separată de kernel. Fiecare pas de antrenare reinterpretează același cod Python. Aceasta funcționează bine până când trebuie să antrenați un model cu 540 de miliarde de parametri pe 2.048 de TPU-uri. Atunci overhead-ul devine prohibitiv.

Notă tehnică a traducerii: Descrierea se referă la execuția eager clasică, nu la întregul PyTorch actual. PyTorch 2 oferă torch.compile, care captează regiuni de graf și generează cod optimizat; nici fiecare expresie nu implică obligatoriu o lansare separată după fuziune. Numerele despre un model de 540 de miliarde de parametri și 2.048 de TPU-uri nu sunt atribuite unei surse în lecție.

Google DeepMind antrenează Gemini pe JAX. Anthropic a antrenat Claude pe JAX. Acestea nu sunt operațiuni mici — sunt unele dintre cele mai mari antrenări de rețele neuronale din lume. Au ales JAX deoarece tratează bucla de antrenare ca pe un program compilabil, nu ca pe o succesiune de apeluri Python.

Notă tehnică a traducerii: Raportul Gemini confirmă folosirea JAX și ML Pathways pe TPU, iar fișa Claude 3 enumeră JAX, PyTorch și Triton drept frameworkuri de bază. Sursele publice nu justifică însă atribuirea exclusivă a alegerii către compilarea buclei și nici superlativul „cele mai mari de pe Pământ”.

JAX este NumPy cu trei superputeri: diferențiere automată, compilare JIT cu XLA și vectorizare automată. Scrieți o funcție care procesează un exemplu. JAX vă oferă o funcție care procesează un lot, calculează gradienți, compilează în cod mașină și rulează pe mai multe dispozitive. Toate acestea fără a schimba funcția originală.

Notă tehnică a traducerii: Aceste capabilități sunt transformări separate, nu sunt aplicate automat toate de JAX. Utilizatorul trebuie să compună explicit grad, jit, vmap și mecanismele de sharding/paralelizare și să furnizeze forme și axe compatibile.

Conceptul

Filosofia JAX

JAX este un framework funcțional. Fără clase, fără stare mutabilă, fără metodă .backward(). În schimb:

PyTorch	JAX
Clasă nn.Module cu stare	Funcție pură: f(params, x) -> y
loss.backward()	jax.grad(loss_fn)(params, x, y)
Execuție eager	Compilare JIT prin XLA
Buclă manuală for x in batch:	Vectorizare automată jax.vmap(f)
DataParallel / FSDP	Paralelism automat jax.pmap(f)
model.parameters() mutabil	Pytree imuabil de array-uri

Aceasta nu este o preferință de stil. Este o constrângere a compilatorului. Compilarea JIT necesită funcții pure — aceleași intrări produc întotdeauna aceleași ieșiri, fără efecte secundare. Această restricție face posibile accelerări de 100 de ori.

Notă tehnică a traducerii: Stilul funcțional facilitează transformările JAX, dar JIT nu oferă un multiplicator universal de 100× și poate accepta efecte controlate sau valori statice cu semantici specifice trasării. Câștigul depinde de operații, forme, hardware, compilare și costul transferurilor.

jax.numpy: suprafața familiară

JAX reimplementează API-ul NumPy pe acceleratoare:

import jax.numpy as jnp

a = jnp.array([1.0, 2.0, 3.0])
b = jnp.array([4.0, 5.0, 6.0])
c = jnp.dot(a, b)

Aceleași nume de funcții. Aceleași reguli de broadcasting. Aceeași semantică a slicing-ului. Însă array-urile se află pe GPU/TPU, iar fiecare operație poate fi urmărită de compilator.

Notă tehnică a traducerii: Array-urile JAX se află pe dispozitivul implicit disponibil, care poate fi CPU, GPU sau TPU; instalarea CPU nu mută datele pe un accelerator. API-ul jax.numpy urmărește NumPy, dar nu este identic în toate detaliile de tipuri, mutabilitate și comportament.

O diferență critică: array-urile JAX sunt imuabile. Nu există a[0] = 5. În schimb: a = a.at[0].set(5). La început pare incomod, apoi devine firesc — imuabilitatea permite compunerea transformărilor precum grad, jit și vmap.

jax.grad: autodiferențiere funcțională

PyTorch atașează gradienți tensorilor (.grad). JAX atașează gradienți funcțiilor.

import jax

def f(x):
    return x ** 2

df = jax.grad(f)
df(3.0)

jax.grad primește o funcție și returnează o funcție nouă care calculează gradientul. Fără apel .backward(). Fără graf de calcul stocat pe tensori. Gradientul este doar o altă funcție pe care o puteți apela, compune sau compila JIT.

Aceasta se compune arbitrar:

d2f = jax.grad(jax.grad(f))
d2f(3.0)

Derivate de ordinul al doilea. De ordinul al treilea. Jacobieni. Hessieni. Toate prin compunerea lui grad. PyTorch poate face și el aceasta (torch.autograd.functional.hessian), dar funcționalitatea a fost adăugată ulterior. În JAX, ea este fundamentul.

Notă tehnică a traducerii: PyTorch oferă în prezent transformări funcționale compozabile prin torch.func, inclusiv grad, vmap, Jacobieni, Hessieni și gradienți pe exemplu. Caracterizarea drept funcționalitate „adăugată artificial” este o opinie istorică, nu o limitare tehnică actuală.

Constrângerea: grad funcționează numai cu funcții pure. Fără instrucțiuni print în interior — rulează în timpul trasării, nu al execuției. Fără mutarea stării externe. Fără generare de numere aleatoare fără gestionarea explicită a cheilor.

jit: compilare cu XLA

@jax.jit
def train_step(params, x, y):
    loss = loss_fn(params, x, y)
    return loss

fast_step = jax.jit(train_step)

La primul apel, JAX trasează funcția — înregistrează operațiile care au loc, fără a executa calculul concret pe date. Apoi transmite reprezentarea către XLA (Accelerated Linear Algebra), compilatorul Google pentru TPU-uri și GPU-uri. XLA fuzionează operații, elimină copii de memorie redundante și generează cod mașină optimizat.

Apelurile ulterioare sar peste corpul Python trasat. Codul compilat rulează pe accelerator la viteza codului nativ.

Când JIT ajută:

• Pași de antrenare — același calcul repetat de mii de ori
• Inferență — același model, intrări diferite
• Orice funcție apelată de mai multe ori cu intrări de forme similare

Când JIT încurcă:

• Funcții cu flux de control Python dependent de valori (if x > 0, unde x este un array urmărit)
• Calcule executate o singură dată — overhead-ul compilării depășește timpul de rulare
• Depanare — trasarea ascunde execuția efectivă

Restricția fluxului de control este reală. jax.lax.cond înlocuiește if/else. jax.lax.scan înlocuiește buclele for. Acestea nu sunt opționale — sunt prețul compilării.

Notă tehnică a traducerii: Funcția Python este executată în timpul trasării cu valori abstracte, iar operațiile concrete sunt apoi compilate; nu este pur și simplu „înregistrată fără execuție”. Fluxul Python static și buclele cu limite statice pot funcționa sub jit; primitivele lax sunt necesare în special pentru control dependent de valori trasate. Recompilarea poate apărea la schimbarea formelor, tipurilor sau argumentelor statice.

vmap: vectorizare automată

Scrieți o funcție care procesează un exemplu:

def predict(params, x):
    return jnp.dot(params['w'], x) + params['b']

vmap o ridică pentru a procesa un lot:

batch_predict = jax.vmap(predict, in_axes=(None, 0))

in_axes=(None, 0) înseamnă: nu creați loturi peste params, care sunt partajați; creați loturi peste axa 0 a lui x. Fără buclă for manuală. Fără reshape. Fără propagarea manuală a dimensiunii lotului. JAX determină dimensiunea lotului și vectorizează întregul calcul.

Aceasta nu este doar conveniență sintactică. vmap generează cod vectorizat fuzionat care rulează de 10–100 de ori mai rapid decât o buclă Python. Și se compune cu jit și grad:

per_example_grads = jax.vmap(jax.grad(loss_fn), in_axes=(None, 0, 0))

Gradienți pe exemplu. O singură linie. În PyTorch este aproape imposibil fără artificii.

Notă tehnică a traducerii: vmap exprimă vectorizarea, dar nu garantează fuziune sau o accelerare de 10–100×; performanța depinde de operații și backend. PyTorch oferă acum echivalentul torch.func.vmap(torch.func.grad(...)) și un tutorial oficial pentru gradienți pe exemplu, deci afirmația „aproape imposibil” este depășită.

pmap: paralelism de date pe mai multe dispozitive

parallel_step = jax.pmap(train_step, axis_name='devices')

pmap replică funcția pe toate dispozitivele disponibile — GPU-uri sau TPU-uri — și împarte lotul. În interiorul funcției, jax.lax.pmean și jax.lax.psum sincronizează gradienții între dispozitive.

Google antrenează Gemini pe mii de cipuri TPU v5e folosind pmap și succesorul său shard_map. Modelul de programare: scrieți versiunea pentru un singur dispozitiv, înveliți-o cu pmap și ați terminat.

Notă tehnică a traducerii: pmap mapează explicit o axă ale cărei dimensiuni trebuie să fie compatibile cu numărul dispozitivelor; nu împarte automat un lot arbitrar. Documentația curentă îl numește metoda veche și recomandă de regulă shard_map. Raportul Gemini confirmă JAX/Pathways și TPU, dar nu susține această descriere simplificată a strategiei concrete prin pmap.

Pytrees: structura universală de date

JAX operează pe „pytrees” — combinații imbricate de liste, tuple, dicționare și array-uri. Parametrii modelului sunt un pytree:

params = {
    'layer1': {'w': jnp.zeros((784, 256)), 'b': jnp.zeros(256)},
    'layer2': {'w': jnp.zeros((256, 128)), 'b': jnp.zeros(128)},
    'layer3': {'w': jnp.zeros((128, 10)),  'b': jnp.zeros(10)},
}

Fiecare transformare JAX — grad, jit, vmap — știe să traverseze pytrees. jax.tree.map(f, tree) aplică f fiecărei frunze. Astfel actualizează optimizatorii toți parametrii simultan:

params = jax.tree.map(lambda p, g: p - lr * g, params, grads)

Fără metodă .parameters(). Fără înregistrarea parametrilor. Structura arborelui este modelul.

Funcțional sau orientat pe obiecte

PyTorch păstrează starea în obiecte:

class Model(nn.Module):
    def __init__(self):
        self.linear = nn.Linear(784, 10)

    def forward(self, x):
        return self.linear(x)

JAX folosește funcții pure cu stare explicită:

def predict(params, x):
    return jnp.dot(x, params['w']) + params['b']

Parametrii sunt transmiși ca argumente. Nimic nu este stocat. Nimic nu este modificat. Astfel, fiecare funcție este testabilă, compozabilă și compilabilă. Înseamnă și că gestionați singuri parametrii — sau folosiți o bibliotecă precum Flax ori Equinox.

Ecosistemul JAX

JAX vă oferă primitive. Bibliotecile oferă ergonomie:

Bibliotecă	Rol	Stil
Flax (Google)	Straturi de rețea neuronală	nn.Module cu stare explicită
Equinox (Patrick Kidger)	Straturi de rețea neuronală	Bazat pe pytree, Pythonic
Optax (DeepMind)	Optimizatori + scheme LR	Transformări compozabile ale gradientului
Orbax (Google)	Puncte de control	Salvare/restaurare pytrees
CLU (Google)	Metrici + jurnalizare	Utilitare pentru bucla de antrenare

Optax este biblioteca standard de optimizare. Separă transformarea gradientului — Adam, SGD, clipping — de actualizarea parametrilor, ceea ce facilitează compunerea:

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adam(learning_rate=1e-3),
)

Când să folosiți JAX sau PyTorch

Factor	JAX	PyTorch
Suport TPU	De prim rang — Google le-a construit pe ambele	Întreținut de comunitate — torch_xla
Suport GPU	Bun — CUDA prin XLA	De prim rang — CUDA nativ
Depanare	Dificilă — trasare + compilare	Ușoară — eager, linie cu linie
Ecosistem	Orientat spre cercetare — Flax, Equinox	Enorm — HuggingFace, torchvision etc.
Angajare	Nișă — Google/DeepMind/Anthropic	Răspândit — pretutindeni
Antrenare la scară mare	Superior — XLA, pmap, mesh	Bun — FSDP, DeepSpeed
Viteza prototipării	Mai lentă — overhead funcțional	Mai rapidă — modificați și continuați
Inferență în producție	TensorFlow Serving, Vertex AI	TorchServe, Triton, ONNX
Cine îl folosește	DeepMind — Gemini; Anthropic — Claude	Meta — Llama; OpenAI — GPT; Stability AI

Notă tehnică a traducerii: Tabelul conține evaluări generale, nu benchmarkuri universale. Ambele ecosisteme acceptă execuție eager și compilată, mai multe backenduri și strategii distribuite; depanarea, prototiparea și performanța depind de program și echipă. torch_xla este un proiect oficial al ecosistemului PyTorch, iar asocierea unor familii de modele cu un singur framework nu descrie neapărat întregul lor stack.

Răspunsul sincer: folosiți PyTorch dacă nu aveți un motiv concret pentru JAX. Aceste motive sunt accesul la TPU, necesitatea gradienților pe exemplu, antrenarea multi-dispozitiv la scară foarte mare sau activitatea la Google/DeepMind/Anthropic.

Numere aleatoare în JAX

JAX nu are o stare aleatoare globală. Fiecare operație aleatoare necesită o cheie PRNG explicită:

key = jax.random.PRNGKey(42)
key1, key2 = jax.random.split(key)
w = jax.random.normal(key1, shape=(784, 256))

La început este incomod. Dar garantează reproductibilitatea între dispozitive și compilări — proprietate pe care torch.manual_seed din PyTorch nu o poate garanta în configurații multi-GPU.

Notă tehnică a traducerii: Cheile explicite fac fluxul aleator reproductibil și controlabil, dar nu oferă o garanție universală între toate backendurile, shardingurile și implementările PRNG. Documentația JAX arată că unele PRNG-uri diferă între CPU/GPU/TPU sau sub vmap; PyTorch documentează la rândul său opțiuni și limite de reproductibilitate, nu o imposibilitate absolută.

batchnorm-effect

Notă tehnică a traducerii: Figure-blocul batchnorm-effect este păstrat literal din sursă, dar lecția nu discută BatchNorm și nu există un asset asociat în directorul original. Nu a fost adăugată nicio imagine aproximativă.

Construiți-l

Pasul 1: configurarea și datele

Vom antrena un MLP cu trei straturi pe MNIST folosind JAX și Optax: 784 de intrări, două straturi ascunse cu 256 și 128 de neuroni și 10 clase de ieșire.

import jax
import jax.numpy as jnp
from jax import random
import optax

def get_mnist_data():
    from sklearn.datasets import fetch_openml
    mnist = fetch_openml('mnist_784', version=1, as_frame=False, parser='auto')
    X = mnist.data.astype('float32') / 255.0
    y = mnist.target.astype('int')
    X_train, X_test = X[:60000], X[60000:]
    y_train, y_test = y[:60000], y[60000:]
    return X_train, y_train, X_test, y_test

Notă tehnică a traducerii: Exemplul depinde și de scikit-learn, deși această dependență lipsește din comenzile de instalare de mai jos. fetch_openml descarcă datele prin rețea și împărțirea prin poziții presupune ordinea standard MNIST returnată de această versiune.

Pasul 2: inițializarea parametrilor

Fără clasă. Doar o funcție care returnează un pytree:

def init_params(key):
    k1, k2, k3 = random.split(key, 3)
    scale1 = jnp.sqrt(2.0 / 784)
    scale2 = jnp.sqrt(2.0 / 256)
    scale3 = jnp.sqrt(2.0 / 128)
    params = {
        'layer1': {
            'w': scale1 * random.normal(k1, (784, 256)),
            'b': jnp.zeros(256),
        },
        'layer2': {
            'w': scale2 * random.normal(k2, (256, 128)),
            'b': jnp.zeros(128),
        },
        'layer3': {
            'w': scale3 * random.normal(k3, (128, 10)),
            'b': jnp.zeros(10),
        },
    }
    return params

Inițializare He, realizată manual. Trei chei PRNG separate dintr-un singur seed. Fiecare pondere este un array imuabil într-un dicționar imbricat.

Notă tehnică a traducerii: Scalarea He se potrivește straturilor ascunse urmate de ReLU, dar stratul final produce logits și nu este urmat de ReLU; pentru el se folosește frecvent o inițializare diferită, de exemplu Xavier/Glorot. În bucla de mai jos, aceeași cheie rădăcină este trecută la init_params(key) și apoi este despărțită pentru amestecare. Pentru un flux aleator explicit, separați mai întâi key, init_key = random.split(key) și apelați init_params(init_key); codul-sursă rămâne păstrat exact.

Pasul 3: propagarea înainte

def forward(params, x):
    x = jnp.dot(x, params['layer1']['w']) + params['layer1']['b']
    x = jax.nn.relu(x)
    x = jnp.dot(x, params['layer2']['w']) + params['layer2']['b']
    x = jax.nn.relu(x)
    x = jnp.dot(x, params['layer3']['w']) + params['layer3']['b']
    return x

def loss_fn(params, x, y):
    logits = forward(params, x)
    one_hot = jax.nn.one_hot(y, 10)
    return -jnp.mean(jnp.sum(jax.nn.log_softmax(logits) * one_hot, axis=-1))

Funcții pure. Parametrii intră, predicția iese. Fără self, fără stare stocată. loss_fn calculează entropia încrucișată de la zero — softmax, logaritm, medie negativă.

Notă tehnică a traducerii: Codul folosește direct log_softmax, o formulare stabilă care combină logaritmul și softmax; nu calculează mai întâi un softmax separat, așa cum sugerează descrierea.

Pasul 4: pas de antrenare compilat JIT

@jax.jit
def train_step(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

@jax.jit
def accuracy(params, x, y):
    logits = forward(params, x)
    preds = jnp.argmax(logits, axis=-1)
    return jnp.mean(preds == y)

jax.value_and_grad returnează valoarea pierderii și gradienții într-o singură transformare. Decoratorul @jax.jit compilează ambele funcții prin XLA. După primul apel și cât timp semnătura compilată rămâne aplicabilă, fiecare pas reutilizează executabilul din cache.

Pasul 5: bucla de antrenare

optimizer = optax.adam(learning_rate=1e-3)

X_train, y_train, X_test, y_test = get_mnist_data()
X_train, X_test = jnp.array(X_train), jnp.array(X_test)
y_train, y_test = jnp.array(y_train), jnp.array(y_test)

key = random.PRNGKey(0)
params = init_params(key)
opt_state = optimizer.init(params)

batch_size = 128
n_epochs = 10

for epoch in range(n_epochs):
    key, subkey = random.split(key)
    perm = random.permutation(subkey, len(X_train))
    X_shuffled = X_train[perm]
    y_shuffled = y_train[perm]

    epoch_loss = 0.0
    n_batches = len(X_train) // batch_size
    for i in range(n_batches):
        start = i * batch_size
        xb = X_shuffled[start:start + batch_size]
        yb = y_shuffled[start:start + batch_size]
        params, opt_state, loss = train_step(params, opt_state, xb, yb)
        epoch_loss += loss

    train_acc = accuracy(params, X_train[:5000], y_train[:5000])
    test_acc = accuracy(params, X_test, y_test)
    print(f"Epoch {epoch + 1:2d} | Loss: {epoch_loss / n_batches:.4f} | "
          f"Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")

Zece epoci. Aproximativ 97% acuratețe pe test. Prima epocă este lentă din cauza compilării JIT. Epocile 2–10 sunt rapide.

Notă tehnică a traducerii: Valoarea de ~97% nu este garantată: depinde de versiunile bibliotecilor, backend, precizie și date. Bucla ignoră ultimele 60000 % 128 = 96 de exemple la fiecare epocă, deoarece folosește împărțire întreagă. Pentru cronometrare corectă pe acceleratoare trebuie ținut cont și de execuția asincronă, de exemplu prin block_until_ready().

Notă tehnică a traducerii: Bucla evaluează setul de test după fiecare epocă, iar primul exercițiu cere tot comparații pe test. Pentru monitorizare, alegerea modelului și hiperparametri, folosiți un set de validare separat și păstrați testul pentru evaluarea finală. Eticheta Train Acc măsoară aici doar primele 5.000 de exemple de antrenare, nu întregul set.

Observați ce lipsește: fără .zero_grad(), fără .backward(), fără .step(). Întreaga actualizare este un singur apel de funcție compusă. Gradienții sunt calculați, transformați de Adam și aplicați parametrilor — toate în interiorul train_step.

Folosiți-l

Flax: standardul Google

Flax este cea mai răspândită bibliotecă JAX pentru rețele neuronale. Readuce nn.Module, dar cu gestionarea explicită a stării:

import flax.linen as nn

class MLP(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(256)(x)
        x = nn.relu(x)
        x = nn.Dense(128)(x)
        x = nn.relu(x)
        x = nn.Dense(10)(x)
        return x

model = MLP()
params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 784)))
logits = model.apply(params, x_batch)

Aceeași structură ca în PyTorch, dar params este separat de model. model.init() creează parametrii. model.apply(params, x) execută propagarea înainte. Obiectul model nu are stare.

Notă tehnică a traducerii: Fragmentul folosește API-ul Flax Linen. Documentația Flax actuală recomandă noilor utilizatori API-ul NNX, care permite și semantică bazată pe obiecte Python și actualizări de stare; afirmația că obiectul model „nu are stare” descrie tiparul Linen din exemplu, nu întregul Flax actual. x_batch nu este definit în fragment, iar fragmentul Equinox următor apelează de asemenea o variabilă x nedefinită.

Equinox: alternativa Pythonică

Equinox, creat de Patrick Kidger, reprezintă modelele ca pytrees:

import equinox as eqx

model = eqx.nn.MLP(
    in_size=784, out_size=10, width_size=256, depth=2,
    activation=jax.nn.relu, key=jax.random.PRNGKey(0)
)
logits = model(x)

Modelul însuși este un pytree. Nu este necesar .apply(). Parametrii sunt pur și simplu frunzele modelului. Această abordare este mai apropiată de modul de gândire JAX.

Optax: optimizatori compozabili

Optax separă transformarea gradientului de actualizare:

schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0, peak_value=1e-3,
    warmup_steps=1000, decay_steps=50000
)

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(learning_rate=schedule, weight_decay=0.01),
)

Clipping al gradientului, încălzirea ratei de învățare, decay al ponderilor — toate compuse într-un lanț de transformări. Fiecare transformare primește gradienții, îi modifică și îi transmite următoarei. Fără o clasă monolitică de optimizator.

Livrați-l

Instalare:

pip install jax jaxlib optax flax

Pentru suport GPU:

pip install jax[cuda12]

Pentru TPU — Google Cloud:

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Notă tehnică a traducerii: Comenzile sunt păstrate literal din sursă, dar instalarea JAX depinde de platformă și versiune; documentația curentă trebuie consultată pentru extra-ul CUDA/TPU acceptat. Exemplul MNIST mai necesită scikit-learn, iar exemplele ulterioare necesită separat Equinox și Orbax.

Capcane de performanță:

• Primul apel JIT este lent din cauza compilării. Faceți o încălzire înainte de benchmark.
• Evitați buclele Python peste array-uri JAX în interiorul JIT. Folosiți jax.lax.scan sau jax.lax.fori_loop.
• jax.debug.print() funcționează în JIT. print() obișnuit nu funcționează.
• Profilați cu jax.profiler sau TensorBoard. Compilarea XLA poate ascunde blocajele.
• JAX prealocă implicit 75% din memoria GPU. Setați XLA_PYTHON_CLIENT_PREALLOCATE=false pentru dezactivare.

Notă tehnică a traducerii: Instrucțiunile print() obișnuite pot rula în timpul trasării și afișa valori abstracte, nu valorile de execuție; pentru acestea se folosește jax.debug.print(). Dezactivarea prealocării poate reduce memoria rezervată, dar poate crește fragmentarea și nu este o optimizare universal recomandată.

Puncte de control:

import orbax.checkpoint as ocp
checkpointer = ocp.PyTreeCheckpointer()
checkpointer.save('/tmp/model', params)
restored = checkpointer.restore('/tmp/model')

Notă tehnică a traducerii: PyTreeCheckpointer aparține API-ului Orbax V0, considerat vechi de documentația actuală. Pentru cod nou, consultați API-ul V1 și StandardCheckpointer; restaurarea poate necesita un țintă abstractă și informații de sharding. Calea de salvare trebuie de asemenea administrată explicit. Blocul originar este păstrat exact.

Această lecție produce:

• outputs/prompt-jax-optimizer.md — un prompt pentru alegerea configurației potrivite a optimizatorului JAX
• outputs/skill-jax-patterns.md — o abilitate care acoperă tiparele funcționale din JAX

Exerciții

1. Adăugați dropout în MLP. În JAX, dropout necesită o cheie PRNG — transmiteți o cheie prin propagarea înainte și separați-o pentru fiecare strat dropout. Comparați acuratețea pe test cu și fără dropout.

2. Folosiți jax.vmap pentru a calcula gradienții pe exemplu pentru un lot de 32 de imagini MNIST. Calculați norma gradientului pentru fiecare exemplu. Care exemple au cei mai mari gradienți și de ce?

3. Înlocuiți funcția manuală de propagare înainte cu una generică, mlp_forward(params, x), care funcționează pentru orice număr de straturi. Folosiți jax.tree.leaves pentru a determina automat adâncimea.

Notă tehnică a traducerii: jax.tree.leaves returnează frunzele într-o ordine determinată de structura pytree, dar nu codifică singur noțiunea de strat și include separat ponderile și biasurile. Pentru o adâncime robustă, păstrați o structură explicită a straturilor sau traversați cheile validate.

4. Măsurați performanța pasului de antrenare cu și fără @jax.jit. Cronometrați câte 100 de pași. Cât de mare este accelerarea pe hardware-ul dumneavoastră? Care este overhead-ul compilării la primul apel?

5. Implementați clipping-ul gradientului compunând optax.chain(optax.clip_by_global_norm(1.0), optax.adam(1e-3)). Antrenați cu și fără clipping. Reprezentați grafic norma gradientului pe parcursul antrenării pentru a observa efectul.

Termeni-cheie

Termen	Ce spun oamenii	Ce înseamnă de fapt
XLA	„Lucrul care face JAX rapid”	Accelerated Linear Algebra — un ecosistem de compilatoare care fuzionează operații și generează kerneluri GPU/TPU optimizate dintr-un graf de calcul
JIT	„Compilare just-in-time”	JAX trasează funcția la primul apel pentru o semnătură, compilează și reutilizează versiunea compilată la apelurile compatibile ulterioare
Funcție pură	„Fără efecte secundare”	O funcție a cărei ieșire depinde numai de intrări — fără stare globală, mutație sau aleatorietate fără chei explicite
vmap	„Auto-batching”	Transformă o funcție pentru un exemplu într-una mapată pe un lot, fără rescriere manuală
pmap	„Paralelism automat”	Replică o funcție pe mai multe dispozitive și mapează o axă a intrării
Pytree	„Dicționar imbricat de array-uri”	Orice structură imbricată de liste, tuple, dicționare și frunze pe care JAX o poate traversa și transforma
Trasare	„Înregistrarea calculului”	JAX execută funcția cu valori abstracte pentru a construi o reprezentare a calculului, fără a calcula rezultatele concrete
Autodiferențiere funcțională	„Gradientul unei funcții”	Calcularea derivatelor prin transformarea funcțiilor, nu prin atașarea stocării gradientului la tensori
Optax	„Biblioteca de optimizare JAX”	Bibliotecă de transformări compozabile ale gradientului — Adam, SGD, clipping și planificare — care pot fi înlănțuite
Flax	„nn.Module pentru JAX”	Bibliotecă de rețele neuronale pentru JAX, care adaugă abstracții de straturi și gestionarea stării

Lecturi suplimentare

• Documentația JAX: https://jax.readthedocs.io/ — documentația oficială, cu tutoriale excelente despre grad, jit și vmap
• „JAX: composable transformations of Python+NumPy programs” (Bradbury și colaboratorii, 2018) — lucrarea originală care explică filosofia de proiectare
• Documentația Flax: https://flax.readthedocs.io/ — biblioteca Google de rețele neuronale pentru JAX
• Patrick Kidger, „Equinox: neural networks in JAX via callable PyTrees and filtered transformations” (2021) — alternativa Pythonică la Flax
• DeepMind, „Optax: composable gradient transformation and optimisation” — biblioteca standard de optimizare
• „You Don’t Know JAX” (Colin Raffel, 2019) — un ghid practic despre capcanele și tiparele JAX, scris de unul dintre autorii T5

Notă tehnică a traducerii: Postarea „You Don’t Know JAX” a fost publicată în ianuarie 2019, nu în 2020 cum indică sursa.

Sursă: Originalul în limba engleză

Navigare: ← Lecția 03.11 — Introducere în PyTorch · Faza 3 — Fundamentele învățării profunde · Catalog complet · Lecția 03.13 — Depanarea rețelelor neuronale →