JAX

Warnings

• breaking bare pip install jax installs a minimal stub — no jaxlib binary. You must use an extra: pip install jax[cpu] for CPU, pip install jax[cuda12] for CUDA 12. Without an extra, import jax raises ImportError or runs extremely slowly.
  Fix: Always install with an extra: pip install 'jax[cpu]' or pip install 'jax[cuda12]'. Check https://jax.readthedocs.io/en/latest/installation.html for the current CUDA extras.
• breaking jax and jaxlib versions must exactly match. Installing mismatched versions raises RuntimeError on import. The JAX team periodically yanks old jaxlib wheels from PyPI.
  Fix: Always install together: pip install 'jax[cpu]' — this installs the matching jaxlib automatically. If pinning: pin both jax==X.Y.Z and jaxlib==X.Y.Z to the same version.
• breaking jax.jit() and other transforms now enforce keyword-only arguments (since 0.7). jax.jit(f, (0,)) for static_argnums raises TypeError. Was DeprecationWarning in 0.6.
  Fix: Pass all transform arguments by keyword: jax.jit(f, static_argnums=(0,)) not jax.jit(f, (0,)).
• breaking jax.pmap is in maintenance mode. New code should use jax.shard_map for multi-device parallelism. pmap's default implementation is being switched to shard_map internals.
  Fix: Migrate new multi-device code to jax.shard_map. Existing pmap code will continue to work for now but will not receive new features.
• breaking Older jaxlib wheels are periodically deleted from PyPI due to storage limits. Pinning old jaxlib versions will cause install failures in CI after deletion.
  Fix: Install older versions from the JAX archive index: pip install 'jax[cpu]==X.Y.Z' -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
• gotcha Side effects inside jit/grad/vmap are silently dropped. Python print(), list.append(), global variable mutations only execute during the initial tracing pass — not on subsequent JIT-compiled calls. The #1 invisible bug for JAX beginners.
  Fix: Use jax.debug.print() for debugging inside jit. Keep all side effects outside of transformed functions. Functions must be pure (same inputs → same outputs).
• gotcha JAX arrays are immutable — no in-place operations. x[0] = 1 raises TypeError. Use x.at[0].set(1) for functional updates.
  Fix: Use the .at[].set() / .at[].add() / .at[].mul() functional update API: x = x.at[0].set(1)
• gotcha By default JAX uses 32-bit floats even when NumPy would use 64-bit. jnp.array(1.0).dtype is float32, not float64. Enable x64 mode explicitly if needed.
  Fix: Enable 64-bit: jax.config.update('jax_enable_x64', True) — must be called before any JAX operations. Or use context manager: with jax.enable_x64(): ...

Install

• pip install jax[cpu] CPU (recommended default)
• pip install jax[cuda12] CUDA 12 GPU
• pip install jax[cuda13] CUDA 13 GPU
• pip install jax[tpu] Google TPU

Imports

• jax.jit

  import jax
  import jax.numpy as jnp
  
  @jax.jit
  def f(x, y):
      return jnp.dot(x, y)
  
  # Or explicit keyword args:
  jax.jit(f, static_argnums=(0,))

  Since 0.7, jax.jit() requires fun to be passed by position and all other arguments (static_argnums, static_argnames, etc.) by keyword. Positional use of these args now raises an error.
• pure functions

  # Pure function — same inputs always give same outputs
  @jax.jit
  def add(x, y):
      return x + y

  JAX requires pure functions. Side effects (Python print, list append, global mutation) inside jit/grad/vmap are silently dropped or only execute during tracing. This is the #1 silent bug for JAX beginners.

Quickstart

Core JAX pattern: compose grad, jit, vmap freely. All functions must be pure.

import jax
import jax.numpy as jnp

# grad: automatic differentiation
def loss(params, x):
    return jnp.sum((params['w'] @ x - params['b']) ** 2)

grad_loss = jax.grad(loss)  # gradient w.r.t. first arg by default

# jit: XLA compilation
@jax.jit
def fast_loss(params, x):
    return loss(params, x)

# vmap: auto-vectorize over batch dimension
batched_loss = jax.vmap(loss, in_axes=(None, 0))  # params fixed, x batched

# Compose freely:
fast_batched_grad = jax.jit(jax.vmap(jax.grad(loss), in_axes=(None, 0)))