JAX gotchas cheat sheet

A quick-reference list of the ways JAX will surprise you when you arrive from NumPy or Fortran. Chapter 1.1 walks through most of these with examples; this page is the condensed version to keep open while you work. See also the official JAX sharp bits page.

Gotcha	Symptom	Fix	Shown in
Arrays are immutable	TypeError on x[i] = v	x = x.at[i].set(v)	1.1
float32 by default	Gradients disagree with finite differences; drifting budgets	jax.config.update("jax_enable_x64", True) before creating arrays	3.5
Python if on traced values	TracerBoolConversionError under jit	jnp.where, lax.cond, lax.select	1.1
Python for loops under jit	Minutes-long compile times	lax.scan (time stepping), lax.fori_loop	1.1
Side effects in jitted functions	print fires once (at trace time), not per call	jax.debug.print; keep functions pure	1.1, 3.5
Implicit global RNG	No np.random.seed equivalent	Explicit keys: jax.random.key, jax.random.split	1.1
Recompilation	First call after a shape/dtype change is slow	Keep shapes static; pad instead of resizing	1.1
nan gradients from safe-looking code	jnp.where with a non-differentiable untaken branch	The “double-where” idiom; jax_debug_nans	3.3, 3.5
Thresholds/switches in physics	Loss varies but gradient is exactly zero	Smooth, reformulate, or custom_vjp	3.3
Long-rollout grad runs out of memory	Adjoint memory grows with steps	jax.checkpoint, nested-scan checkpointing	3.2
Integer/weak-type promotion	Silent dtype surprises in mixed expressions	Be explicit with dtypes in model constants	—

(Starter list — group members: add gotchas as you hit them, one row each.)