G1Joshi/Agent-Skills

JAX

JAX es "NumPy con esteroides". Combina Autograd (diferenciación automática) con XLA (compilación). En 2025, Flax NNX (OOP estilo PyTorch) se está convirtiendo en el estándar.

Cuándo Usarlo

• Computación Científica y Numérica de Alto Rendimiento: Compilación Autograd y XLA para álgebra lineal acelerada en GPU/TPU.
• Investigación Personalizada en Deep Learning: Construcción de redes neuronales con Flax, Haiku o Equinox sin la sobrecarga de frameworks.
• Transformaciones de Funciones Componibles: Encadenamiento fluido de jit (compilación), grad (derivadas) y vmap (vectorización).
• Entrenamiento Distribuido en Múltiples Dispositivos: Paralelización de cálculos en clústeres usando pmap y mapas de fragmentos (jax.experimental.shard_map).

Inicio Rápido

import jax
import jax.numpy as jnp

# Multiplicación de matrices compilada con JIT y diferenciación automática
@jax.jit
def loss_fn(w, x, y):
    pred = jnp.dot(x, w)
    return jnp.mean((pred - y) ** 2)

# Calcular valor y gradientes simultáneamente
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("Gradientes:", grads)

Conceptos Clave

#Transformaciones de Funciones Componibles: jit, grad y vmap

Combinando compilación justo-a-tiempo, diferenciación automática y procesamiento por lotes automatizado:

import jax
import jax.numpy as jnp

# Definición de función pura
def loss_fn(w, x, y):
    pred = jnp.dot(x, w)
    return jnp.mean((pred - y) ** 2)

# Transformación 1: Gradiente automático con respecto a los pesos
grad_fn = jax.grad(loss_fn)

# Transformación 2: Compilación JIT mediante XLA para aceleración extrema en GPU
fast_grad_fn = jax.jit(grad_fn)

# Transformación 3: Vectorización automática entre lotes
# batch_loss_fn maneja matrices 2D automáticamente sin bucles
vmapped_loss = jax.vmap(loss_fn, in_axes=(None, 0, 0))

# Cálculo de ejemplo
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("Gradiente calculado:", grad_val)

#Gestión Explícita de Claves PRNG

Generación de números pseudoaleatorios sin estado:

import jax

# Inicializar clave PRNG raíz
key = jax.random.PRNGKey(2026)

# Las claves JAX deben dividirse explícitamente; nunca reutilizarse
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("Media de muestra normal:", jnp.mean(data_normal))

#Modelo con Estado usando Equinox y Funciones Puras

Redes neuronales expresivas y orientadas a objetos basadas en funciones JAX puras:

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("Salida del modelo:", output)

Patrones Comunes

Vectorización por Lotes con vmap

Problema: Escribir bucles manuales en Python para evaluar funciones sobre lotes ralentiza la ejecución.

Solución: Usar jax.vmap para vectorizar automáticamente funciones de un solo ejemplo:

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

# Vectorizar automáticamente a través de la dimensión del lote (eje 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) # Forma de salida: (100,)

Mejores Prácticas (2026)

• Haz escribir código puramente funcional: sin efectos secundarios, sin mutaciones in-place de arrays (x.at[idx].set(val) en lugar de x[idx] = val).
• Haz envolver funciones críticas para el rendimiento con @jax.jit para activar la compilación XLA.
• Haz dividir explícitamente las claves PRNG (jax.random.split(key)) cada vez que se generen números aleatorios.
• Haz usar bibliotecas modernas de alto nivel como Equinox o Flax Linen en lugar de escribir diccionarios de parámetros manuales.
• No uses condicionales estándar de Python (if x > 0:) dentro de funciones JIT sobre trazadores dinámicos; usa jax.lax.cond.
• No reutilices claves PRNG; reutilizarlas genera números aleatorios correlacionados estadísticamente.
• No mutes variables globales dentro de funciones transformadas con jit, grad o vmap.

Solución de Problemas

Error	Causa	Solución
ConcretizationTypeError: Abstract tracer value used in Python if/while	Ramificación condicional de Python depende de valores dinámicos de arrays JAX dentro de jit.	Usar jax.lax.cond o jax.lax.while_loop en lugar de if estándar de Python.
JAX arrays are immutable	Se intentó una asignación in-place (arr[0] = 5).	Usar sintaxis de actualización funcional: arr = arr.at[0].set(5).
PRNGKey reuse warning / identical randomness	Reutilizar la misma clave PRNG produce números aleatorios idénticos.	Dividir claves antes de cada operación aleatoria: key, subkey = jax.random.split(key).

Referencias

• Documentación de JAX