G1Joshi/Agent-Skills

JAX

JAX ist „NumPy auf Steroiden“. Es vereint Autograd (automatische Differenzierung) mit XLA (Kompilierung). Im Jahr 2025 wird Flax NNX (PyTorch-ähnliche OOP) zum Standard.

Verwendung

• Hochleistungsfähiges numerisches und wissenschaftliches Rechnen: Autograd und XLA-Kompilierung für beschleunigte lineare Algebra auf GPU/TPU.
• Eigene Deep-Learning-Forschung: Aufbau neuronaler Netze mit Flax, Haiku oder Equinox ohne Framework-Overhead.
• Zusammensetzbare Funktionstransformationen: Nahtloses Verketten von jit (Kompilierung), grad (Ableitungen) und vmap (Vektorisierung).
• Multi-Device-verteiltes Training: Parallelisierung von Berechnungen über Cluster mit pmap und Shard-Maps (jax.experimental.shard_map).

Schnellstart

import jax
import jax.numpy as jnp

# JIT-kompilierte Matrixmultiplikation und automatische Differenzierung
@jax.jit
def loss_fn(w, x, y):
    pred = jnp.dot(x, w)
    return jnp.mean((pred - y) ** 2)

# Gleichzeitige Berechnung von Wert und Gradienten
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("Gradienten:", grads)

Kernkonzepte

#Zusammensetzbare Funktionstransformationen: jit, grad & vmap

Kombination von Just-in-Time-Kompilierung, automatischer Differenzierung und automatisierter Batch-Verarbeitung:

import jax
import jax.numpy as jnp

# Definition einer reinen Funktion
def loss_fn(w, x, y):
    pred = jnp.dot(x, w)
    return jnp.mean((pred - y) ** 2)

# Transformation 1: Automatischer Gradient bezüglich der Gewichte
grad_fn = jax.grad(loss_fn)

# Transformation 2: JIT-Kompilierung via XLA für extreme GPU-Beschleunigung
fast_grad_fn = jax.jit(grad_fn)

# Transformation 3: Automatische Vektorisierung über Batches
# batch_loss_fn verarbeitet 2D-Matrizen automatisch ohne Schleifen
vmapped_loss = jax.vmap(loss_fn, in_axes=(None, 0, 0))

# Beispielberechnung
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("Berechneter Gradient:", grad_val)

#Explizites PRNG-Key-Management

Zustandsfreie Pseudozufallszahlengenerierung:

import jax

# Initialer PRNG-Key
key = jax.random.PRNGKey(2026)

# JAX-Keys müssen explizit geteilt werden; nie wiederverwenden
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("Mittelwert der Normalverteilung:", jnp.mean(data_normal))

#Zustandsbehaftetes Modell mit Equinox & reinen Funktionen

Ausdrucksstarke, objektorientierte neuronale Netze, die auf reinen JAX-Funktionen basieren:

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("Modellausgabe:", output)

Häufige Muster

Batch-Vektorisierung mit vmap

Problem: Manuelle Python-Schleifen zur Auswertung von Funktionen über Batches verlangsamen die Ausführung.

Lösung: Verwenden Sie jax.vmap, um Einzelbeispielfunktionen automatisch zu vektorisieren:

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

# Automatische Vektorisierung über die Batch-Dimension (Achse 0 von 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) # Ausgabeform: (100,)

Best Practices (2026)

• Tun Sie: Schreiben Sie rein funktionalen Code: keine Seiteneffekte, keine In-Place-Mutationen von Arrays (x.at[idx].set(val) statt x[idx] = val).
• Tun Sie: Umwickeln Sie leistungskritische Funktionen mit @jax.jit, um XLA-Kompilierung auszulösen.
• Tun Sie: Teilen Sie PRNG-Keys explizit (jax.random.split(key)) jedes Mal, wenn Zufallszahlen generiert werden.
• Tun Sie: Verwenden Sie moderne High-Level-Bibliotheken wie Equinox oder Flax Linen, anstatt rohe Parameterdictionaries zu schreiben.
• Nicht tun: Verwenden Sie keine standardmäßigen Python-Bedingungen (if x > 0:) innerhalb von JIT-Funktionen für dynamische Tracer; verwenden Sie jax.lax.cond.
• Nicht tun: Verwenden Sie PRNG-Keys nicht wieder; die Wiederverwendung erzeugt statistisch korrelierte Zufallszahlen.
• Nicht tun: Mutieren Sie keine globalen Variablen innerhalb von Funktionen, die mit jit, grad oder vmap transformiert wurden.

Fehlerbehebung

Fehler	Ursache	Lösung
ConcretizationTypeError: Abstract tracer value used in Python if/while	Python-bedingte Verzweigung hängt von dynamischen JAX-Array-Werten in jit ab.	Verwenden Sie jax.lax.cond oder jax.lax.while_loop anstelle von Python if.
JAX arrays are immutable	In-Place-Zuweisung versucht (arr[0] = 5).	Verwenden Sie funktionale Updatesyntax: arr = arr.at[0].set(5).
PRNGKey reuse warning / identical randomness	Wiederverwendung identischer PRNG-Keys erzeugt identische Zufallszahlen.	Teilen Sie Keys vor jeder Zufallsoperation: key, subkey = jax.random.split(key).

Referenzen

• JAX-Dokumentation