JAX

Common errors

Warnings

Install

Install compatibility stale last tested: 2026-05-12

Imports

• jax.jit
  wrong

  # jit now requires fun as positional, all other args as keyword
  # jax.jit(f, (0,))  — positional static_argnums raises DeprecationWarning in 0.6, error in 0.7+

  correct

  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
  wrong

  # Side effects are silently dropped under jit
  global_list = []
  
  @jax.jit
  def append_and_return(x):
      global_list.append(x)  # silently does nothing under jit
      return x * 2

  correct

  # 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 draft last tested: 2026-04-23

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)))