Advanced

JAX Best Practices

Avoid common pitfalls, debug effectively, write performant JAX code, and follow proven patterns for production-quality JAX applications.

Common Pitfalls and How to Avoid Them

  1. Arrays are immutable

    Never try x[0] = 5. Use x = x.at[0].set(5) instead. This creates a new array (but JAX optimizes it under JIT).

  2. JIT tracing vs execution

    Python side effects (print, list.append) only run during tracing, not on subsequent JIT calls. Use jax.debug.print() for debugging inside JIT.

  3. Random number handling

    Always split PRNG keys explicitly. Never reuse a key: key, subkey = jax.random.split(key).

  4. Control flow inside JIT

    Use jax.lax.cond instead of Python if/else, and jax.lax.fori_loop instead of Python for loops when the condition depends on array values.

Debugging JAX Code

Python
import jax

# Debug prints inside JIT
@jax.jit
def train_step(params, x, y):
    loss = loss_fn(params, x, y)
    jax.debug.print("loss: {}", loss)  # Works inside JIT!
    return loss

# Disable JIT for debugging
with jax.disable_jit():
    result = train_step(params, x, y)  # Runs as pure Python

# Check for NaN/Inf
jax.config.update("jax_debug_nans", True)  # Raises error on NaN

# Print compilation info
jax.config.update("jax_log_compiles", True)

Performance Tips

TipWhy It Matters
JIT everythingUn-JITted code dispatches each op individually to the GPU - massive overhead
Avoid Python loops over dataUse vmap or jax.lax.scan instead of for loops
Minimize host-device transfersKeep data on the device; avoid .numpy() in hot loops
Use donate_argnumsjax.jit(fn, donate_argnums=(0,)) lets JAX reuse input buffers
Profile with JAX profilerjax.profiler.trace() generates TensorBoard-compatible profiles
Batch operationsLarge batched operations are much more efficient than many small ones

PRNG Key Management

Python
import jax

key = jax.random.PRNGKey(0)

# WRONG: reusing the same key gives identical results
# a = jax.random.normal(key, (3,))
# b = jax.random.normal(key, (3,))  # Same as a!

# RIGHT: split the key for each use
key, k1, k2 = jax.random.split(key, 3)
a = jax.random.normal(k1, (3,))
b = jax.random.normal(k2, (3,))  # Different from a

# Pattern for training loops
for step in range(num_steps):
    key, subkey = jax.random.split(key)
    params, loss = train_step(params, subkey, batch)

When to Use JAX

Choose JAX when: You need maximum performance on TPUs, are doing research requiring per-sample gradients or custom differentiation, want composable transforms (jit + grad + vmap), or are working on large-scale distributed training at Google Cloud scale.

Course Complete!

You now understand JAX from fundamentals to advanced topics. Continue learning by exploring FastAI for rapid prototyping or Hugging Face Transformers for NLP.

Next Course: FastAI →
← Advanced Topics Course Overview →

Ready to Go Deeper?

Live instructor-led courses from our partners. Affiliate disclosure.

🧠

AI & ML Courses - 30% Off

Live instructor-led AI, machine learning, data science, and cloud courses for working professionals. Use code Limited30 at checkout.

Edureka
📊

DataCamp - AI & Data Science

Hands-on Python, machine learning, and AI courses with interactive exercises and real projects.

DataCamp
🎓

edX - Top AI Courses

University-level AI courses from MIT, Harvard, Stanford. Earn certificates that employers recognize.

edX