Chapter 8: The Debugging Toolkit

import os, sys, io, re, gc, time, json, logging, subprocess, textwrap, traceback, contextlib, functools, warnings
from functools import partial
from collections import Counter

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
import jax.monitoring as monitoring
from jax.experimental import checkify
import jax.test_util                      # must be imported explicitly (see Chapter 3); with only `import jax`, accessing jax.test_util raises AttributeError
from jax.test_util import check_grads

print("JAX", jax.__version__, "| backend", jax.default_backend(), "| devices", jax.devices())
CONFIG_AT_START = dict(jax.config.values)   # record the full config at the start, to check at the end that all debug switches were restored


def to64(a):
    """Convert to NumPy float64 for cross-checking."""
    return np.asarray(a, dtype=np.float64)


def show_error(fn):
    """Run fn; on error, print "type: first line of message" and return the exception object; return None if there's no error."""
    try:
        out = fn()
    except Exception as e:
        first = str(e).strip().splitlines()[0] if str(e).strip() else ""
        print(f"  → {type(e).__name__}: {first[:150]}")
        return e
    print("  → no error, result:", out)
    return None


def capture_stdout(fn):
    """Run fn and collect stdout. jax.debug.print is an asynchronous callback, so call effects_barrier() at the end to wait for it to finish printing."""
    buf = io.StringIO()
    with contextlib.redirect_stdout(buf):
        jax.block_until_ready(fn())
        jax.effects_barrier()
    return buf.getvalue()


class capture_jax_logs:
    """Temporarily take over the 'jax' logger and collect log messages in self.messages (instead of printing to stderr)."""
    def __enter__(self):
        self.logger = logging.getLogger("jax")
        self.messages = []
        self.handler = logging.Handler()
        self.handler.emit = lambda record: self.messages.append(record.getMessage())
        self._old_propagate = self.logger.propagate
        self.logger.addHandler(self.handler)
        self.logger.propagate = False
        return self

    def __exit__(self, *exc):
        self.logger.removeHandler(self.handler)
        self.logger.propagate = self._old_propagate
        return False


class CompileCounter:
    """Count XLA backend compilations inside the with block; with name="f", only count jit(f)."""
    EVENT = "/jax/core/compile/backend_compile_duration"

    def __init__(self, name=None):
        self.name, self.events = name, []
        self._listener = self._on_event          # keep the same bound-method object so it can be found when unregistering

    def _on_event(self, event, duration, **kw):
        fun = kw.get("fun_name", "?")
        if event == self.EVENT and (self.name is None or fun == f"jit({self.name})"):
            self.events.append(fun)

    def __enter__(self):
        monitoring.register_event_duration_secs_listener(self._listener)
        return self

    def __exit__(self, *exc):
        monitoring.unregister_event_duration_listener(self._listener)
        return False

    @property
    def count(self):
        return len(self.events)


def run_py(code, *args, env=None, stdin=None, timeout=300):
    """Run code in a fresh Python subprocess and return the CompletedProcess. Keys in env override the current environment variables."""
    full_env = {k: v for k, v in os.environ.items() if not k.startswith(("XLA_FLAGS", "JAX_"))}
    full_env["PYTHON_COLORS"] = "0"            # Python 3.14's colored tracebacks mix ANSI escape codes into the output
    full_env.update(env or {})
    return subprocess.run([sys.executable, "-c", textwrap.dedent(code), *args], input=stdin,
                          capture_output=True, text=True, env=full_env, timeout=timeout)


def result_lines(proc):
    """Extract what follows 'RESULT ' in the subprocess stdout (it may come right after the debugger prompt '(jdb) '); everything else is treated as noise."""
    return [l.split("RESULT ", 1)[1] for l in proc.stdout.splitlines() if "RESULT " in l]

@jax.jit
def stats(x):
    jax.debug.print("mean={m:.3f}  max={mx:.2f}  size={n}", m=x.mean(), mx=x.max(), n=x.size)
    jax.debug.print("positional: first={} last={}", x[0], x[-1])
    return x * 2

out = capture_stdout(lambda: stats(jnp.array([1.0, 2.0, 4.5])))
print(out)
assert "mean=2.500" in out and "max=4.50" in out and "size=3" in out
assert "positional: first=1.0 last=4.5" in out

# ❌ Wrong: the f-string is evaluated at trace time, so it prints a tracer, and every later run prints that same line
@jax.jit
def fstring_print(x):
    jax.debug.print(f"x = {x}")
    return x

try:
    out_bad = capture_stdout(lambda: fstring_print(jnp.float32(3.0)))
    print("Actual output:", out_bad.strip())
except Exception as e:                     # if the tracer's string happens to contain braces, you also get a format error
    out_bad = f"{type(e).__name__}: {e}"
    print("Actual error:", out_bad)
assert "Tracer" in out_bad and "3.0" not in out_bad

# ✅ Correct: the format string is a plain string and the values are passed as arguments
@jax.jit
def good_print(x):
    jax.debug.print("x = {x}", x=x)
    return x

out_good = capture_stdout(lambda: good_print(jnp.float32(3.0)))
print("Actual output:", out_good.strip())
assert out_good.strip() == "x = 3.0"

log_order = jax.debug.print(ordered=True)          # step one: pass only the options, get back a print function

@jax.jit
def stage_a(x):
    for i in range(3):                              # the Python loop unrolls into 3 prints at trace time
        log_order("A{i}", i=i)
    return x + 1

@jax.jit
def stage_b(x):
    log_order("B x={x}", x=x)
    return x * 2

out = capture_stdout(lambda: stage_b(stage_a(jnp.float32(1.0))))
print(out)
assert out.split() == ["A0", "A1", "A2", "B", "x=2.0"]

dbg_logger = logging.getLogger("jax._src.debugging")
old_level = dbg_logger.level
dbg_logger.setLevel(logging.INFO)              # at the default level, INFO isn't emitted
try:
    with capture_jax_logs() as logs:
        jax.block_until_ready(jax.jit(lambda x: (jax.debug.log("loss = {x:.2f}", x=x), x * 2)[1])(jnp.float32(0.25)))
        jax.effects_barrier()
finally:
    dbg_logger.setLevel(old_level)
print("logging received:", logs.messages)
assert logs.messages == ["loss = 0.25"]

def body(x):
    jax.debug.print("saw x = {x}", x=x)
    return jnp.sin(x) * x

def branchy(x):
    return lax.cond(x > 0,
                    lambda x: (jax.debug.print("positive branch {x}", x=x), x)[1],
                    lambda x: (jax.debug.print("negative branch {x}", x=x), -x)[1], x)

out_v = capture_stdout(lambda: jax.vmap(body)(jnp.array([1.0, 2.0, 3.0])))
out_g = capture_stdout(lambda: jax.grad(body)(2.0))
out_jg = capture_stdout(lambda: jax.jit(jax.grad(body))(2.0))
out_c = capture_stdout(lambda: jax.jit(branchy)(jnp.float32(1.0)))
out_vc = capture_stdout(lambda: jax.vmap(branchy)(jnp.array([1.0, -1.0])))
out_vo = capture_stdout(lambda: jax.vmap(lambda x: (jax.debug.print(ordered=True)("ordered {x}", x=x), x)[1])(jnp.arange(2.0)))
for name, o in [("vmap", out_v), ("grad", out_g), ("jit(grad)", out_jg), ("cond", out_c),
                ("vmap(cond)", out_vc), ("vmap + ordered", out_vo)]:
    print(f"{name:15s}", o.strip().replace("\n", " | "))

assert out_v.splitlines() == ["saw x = 1.0", "saw x = 2.0", "saw x = 3.0"]
assert out_g.splitlines() == ["saw x = 2.0"] and out_jg.splitlines() == ["saw x = 2.0"]
assert out_c.splitlines() == ["positive branch 1.0"]
assert sorted(out_vc.splitlines()) == sorted(["positive branch 1.0", "positive branch -1.0", "negative branch 1.0", "negative branch -1.0"])
assert out_vo.splitlines() == ["ordered 0.0", "ordered 1.0"]

GRAD_TAPS = []

@partial(jax.custom_vjp, nondiff_argnums=(1,))
def grad_tap(x, name):
    """Identity function; when backprop passes through here, hand the incoming gradient to the host to record."""
    return x

def _grad_tap_fwd(x, name):
    return x, None

def _grad_tap_bwd(name, _, g):
    jax.debug.callback(lambda g: GRAD_TAPS.append((name, np.asarray(g))), g)
    return (g,)

grad_tap.defvjp(_grad_tap_fwd, _grad_tap_bwd)

def tap_loss(w, x):
    h = grad_tap(jnp.tanh(w * x), "h")
    return jnp.sum(h ** 2)

w_t, x_t = jnp.array([1.0, 2.0]), jnp.array([0.5, 0.1])
g_w = jax.jit(jax.grad(tap_loss))(w_t, x_t)
jax.effects_barrier()
name, g_h = GRAD_TAPS[-1]
print(f"{name}'s incoming gradient:", g_h, "| gradient w.r.t. w:", g_w)
h64 = np.tanh(to64(w_t) * to64(x_t))
assert name == "h" and np.allclose(g_h, 2 * h64, atol=1e-6)                       # ∂L/∂h = 2h
assert np.allclose(g_w, 2 * h64 * (1 - h64 ** 2) * to64(x_t), atol=1e-6)          # the probe doesn't change the gradient

class ActivationLog:
    """Record per-layer activation statistics inside jit using jax.debug.callback."""
    def __init__(self):
        self.records = []

    def record(self, name, x):
        def _host(m, s, frac_dead):              # runs on the host; the arguments are NumPy arrays
            self.records.append({"layer": name, "mean": float(m), "std": float(s), "dead": float(frac_dead)})
        jax.debug.callback(_host, x.mean(), x.std(), jnp.mean(x <= 0))
        return x

act_log = ActivationLog()

def mlp_forward(params, x):
    for i, (w, b) in enumerate(params):
        x = x @ w + b
        if i < len(params) - 1:
            x = act_log.record(f"relu{i}", jax.nn.relu(x))
    return x

k1, k2, k3 = jax.random.split(jax.random.key(0), 3)
params_cb = [(jax.random.normal(k1, (8, 16)) / jnp.sqrt(8.0), jnp.zeros(16)),
             (jax.random.normal(k2, (16, 16)) / 4.0, jnp.zeros(16)),
             (jax.random.normal(k3, (16, 1)) / 4.0, jnp.zeros(1))]
x_cb = jax.random.normal(jax.random.key(1), (32, 8))
jax.block_until_ready(jax.jit(mlp_forward)(params_cb, x_cb))
jax.effects_barrier()
for r in act_log.records:
    print(f"  {r['layer']}: mean={r['mean']:.3f} std={r['std']:.3f} dead fraction={r['dead']:.2f}")

# Cross-check: recompute the first layer's statistics with NumPy in eager mode
h0 = np.maximum(to64(x_cb) @ to64(params_cb[0][0]), 0)
assert [r["layer"] for r in act_log.records] == ["relu0", "relu1"]
assert abs(act_log.records[0]["mean"] - h0.mean()) < 1e-5 and abs(act_log.records[0]["std"] - h0.std()) < 1e-5

# Under vmap the callback runs once per example (4 examples × 2 layers = 8 records)
act_log.records.clear()
jax.block_until_ready(jax.vmap(mlp_forward, in_axes=(None, 0))(params_cb, x_cb[:4]))
jax.effects_barrier()
print("Number of records under vmap:", len(act_log.records))
assert len(act_log.records) == 8

BP_SCRIPT = """
import jax, jax.numpy as jnp

@jax.jit
def f(x):
    y = jnp.sin(x) * 2
    jax.debug.breakpoint()
    return y + 1

print("RESULT", f(jnp.arange(3.0)).tolist())
"""
proc = run_py(BP_SCRIPT, stdin="p y\np x\nc\n")
print(proc.stdout)
assert "Entering jdb" in proc.stdout and "(jdb)" in proc.stdout
assert "1.6829419" in proc.stdout                                   # p y: sin(1)*2, already a concrete value
assert "Array([0., 1., 2.], dtype=float32)" in proc.stdout          # p x
res = json.loads(result_lines(proc)[0])                               # after typing c, the program runs to completion
assert np.allclose(res, np.sin(np.arange(3.0)) * 2 + 1, atol=1e-6)

COND_BP_SCRIPT = """
import jax, jax.numpy as jnp
from jax import lax

def stop_if_nan(x, y):            # pass in the variables you want to inspect as arguments
    jax.debug.breakpoint()

@jax.jit
def f(x):
    y = jnp.log(x)
    lax.cond(jnp.any(jnp.isnan(y)), stop_if_nan, lambda x, y: None, x, y)
    return y

print("RESULT ok-run", f(jnp.array([1.0, 2.0])).tolist())     # doesn't stop
print("RESULT nan-run", f(jnp.array([1.0, -2.0])).tolist())   # stops; continue after inspecting
"""
proc = run_py(COND_BP_SCRIPT, stdin="p y\np x\nc\n")
print(proc.stdout)
assert proc.stdout.count("Entering jdb") == 1                        # only the run with a NaN stops
assert "Array([ 0., nan], dtype=float32)" in proc.stdout and "Array([ 1., -2.], dtype=float32)" in proc.stdout
assert [l.split()[0] for l in result_lines(proc)] == ["ok-run", "nan-run"]

seen = []
with jax.disable_jit():
    @jax.jit
    def jitted_branch(x):
        seen.append(("argument type inside jit", type(x).__name__))
        return x if x > 0 else -x                       # Python if: only legal under disable_jit
    r1 = jitted_branch(jnp.array(-3.0))
    lax.scan(lambda c, x: (seen.append(("carry inside scan body", type(c).__name__)) or (c + x, c)), jnp.float32(0), jnp.arange(2.0))
    lax.fori_loop(0, 2, lambda i, v: (seen.append(("fori index", type(i).__name__)) or v + i), 0)
    jax.vmap(lambda x: seen.append(("argument inside vmap", type(x).__name__)) or x)(jnp.arange(2.0))
for s in seen:
    print("  ", *s)
kinds = dict(seen)
assert float(r1) == 3.0 and kinds["argument type inside jit"] == "ArrayImpl"
assert kinds["fori index"] == "int" and kinds["argument inside vmap"] == "BatchTracer"

def bisect_numeric(fn, *args, ref):
    """Run the same function under jit and under disable_jit, and compare the max error against the reference."""
    err_jit = float(np.abs(to64(jax.jit(fn)(*args)) - ref).max())
    with jax.disable_jit():
        err_eager = float(np.abs(to64(fn(*args)) - ref).max())
    print(f"    jit error = {err_jit:.3g} | disable_jit error = {err_eager:.3g}")
    return err_jit, err_eager

W_bug = jax.random.normal(jax.random.key(1), (128, 128)) / jnp.sqrt(128.0)
X_bug = jax.random.normal(jax.random.key(2), (50, 128))
ref_bug = (to64(W_bug) @ to64(X_bug).T).sum(0)                     # NumPy float64 reference

# ❌ Wrong: the weights are captured as a constant by the jit closure (jaxlib 0.12's CPU backend on macOS / Apple Silicon silently returns wrong results)
def closure_version(X):
    return jax.vmap(lambda x: jnp.sum(W_bug @ x))(X)

try:
    err_jit, err_eager = bisect_numeric(closure_version, X_bug, ref=ref_bug)
except Exception as e:
    print("Actual error:", type(e).__name__, e)
    raise
assert err_eager < 1e-3
import platform, sys
if sys.platform == "darwin" and platform.machine() == "arm64":   # on Apple Silicon it must reproduce; if not, it may have been fixed
    assert err_jit > 1.0, "Not reproduced on macOS arm64: it may be fixed in a newer jaxlib, so update the notes"
if err_jit > 1.0:
    print("jit wrong, eager right → the problem is in the compilation layer")
else:
    print("This machine doesn't trigger the bug, so jit is correct too (it reproduces on macOS / Apple Silicon and did not show up in our Linux x86 test). The bisection idea is the same.")

# ✅ Correct: pass the weights in as a function argument
def param_version(W, X):
    return jax.vmap(lambda x: jnp.sum(W @ x))(X)

err_jit, err_eager = bisect_numeric(param_version, W_bug, X_bug, ref=ref_bug)
assert err_jit < 1e-3 and err_eager < 1e-3

SCALE = 2.0                                       # module-level "config"

def scaled(x):
    return x * SCALE

jaxpr_txt = str(jax.make_jaxpr(scaled)(1.0))
print(jaxpr_txt)
assert "mul a 2.0:f32[]" in jaxpr_txt             # SCALE shows up as a literal: it was baked in

# ❌ Wrong: keep the config in a global variable and change it after jit
scaled_jit = jax.jit(scaled)
before = float(scaled_jit(1.0))
SCALE = 10.0
try:
    after = float(scaled_jit(1.0))
    print(f"before changing SCALE = {before}, after = {after} (no error, but it still uses the old value)")
except Exception as e:
    print("Actual error:", type(e).__name__, e)
    raise
assert before == after == 2.0

# ✅ Correct: pass the config in as an argument; in the jaxpr it becomes the input variable b instead of a literal
def scaled_ok(x, scale):
    return x * scale

print(jax.make_jaxpr(scaled_ok)(1.0, SCALE))
scaled_ok_jit = jax.jit(scaled_ok)
assert float(scaled_ok_jit(1.0, 2.0)) == 2.0 and float(scaled_ok_jit(1.0, 10.0)) == 10.0
SCALE = 2.0

def scoped_mlp(params, x):
    with jax.named_scope("encoder"):
        h = jax.nn.relu(x @ params["w1"])
    with jax.named_scope("head"):
        out = h @ params["w2"]
        loss = jnp.mean(out ** 2)
    return loss

p_ns = {"w1": jnp.ones((4, 8)) * 0.1, "w2": jnp.ones((8, 2)) * 0.1}
x_ns = jnp.ones((3, 4))
jf_ns = jax.jit(scoped_mlp)

# ① jaxpr: the name stack of each equation
traced = jf_ns.trace(p_ns, x_ns)
stacks = [str(e.source_info.name_stack) for e in traced.jaxpr.jaxpr.eqns]
print("① jaxpr name stacks:", stacks)

# ② StableHLO: the default text has no names; only debug_info=True does
lowered = traced.lower()
plain, dbg = lowered.as_text(), lowered.as_text(debug_info=True)
print("② Default StableHLO text contains 'encoder'?", "encoder" in plain, "| with debug_info=True:",
      [l.strip() for l in dbg.splitlines() if "encoder/" in l][:2])

# ③ Optimized HLO: op_name in the metadata
compiled = lowered.compile()
def hlo_ops_in_scope(hlo_text, scope):
    """Find instructions in the optimized HLO whose op_name contains /scope/; return (instruction name, op_name)."""
    pat = re.compile(r"%([\w.\-]+) = .*?op_name=\"([^\"]*/" + re.escape(scope) + r"/[^\"]*)\"")
    return [m.groups() for m in pat.finditer(hlo_text)]
enc_ops = hlo_ops_in_scope(compiled.as_text(), "encoder")
print("③ Instructions in the optimized HLO that belong to encoder:")
for name, op in enc_ops:
    print(f"     %{name:20s} ← {op}")

assert stacks.count("encoder") >= 2 and stacks.count("head") >= 3
assert "encoder" not in plain and "jit(scoped_mlp)/encoder/dot_general" in dbg
assert any(op.endswith("encoder/dot_general") for _, op in enc_ops)
assert np.isclose(float(compiled(p_ns, x_ns)), float(np.mean((np.maximum(to64(x_ns) @ to64(p_ns["w1"]), 0) @ to64(p_ns["w2"])) ** 2)))

print("Types:", type(traced).__name__, "→", type(lowered).__name__, "→", type(compiled).__name__)
print("Traced.out_info:", traced.out_info, "| keys of Compiled.input_shardings[0]:", sorted(compiled.input_shardings[0][0]))
assert isinstance(traced, jax.stages.Traced) and isinstance(lowered, jax.stages.Lowered)
assert isinstance(compiled, jax.stages.Compiled)
assert traced.out_info.shape == () and traced.out_info.dtype == jnp.float32

with jax.debug_nans(True):                          # the with form isn't affected by the vmap missed-detection issue below, so no need to clear caches
    print("inside the block:", jax.config.jax_debug_nans)
    e_nan = show_error(lambda: jax.jit(lambda x: jnp.log(x))(jnp.array(-1.0)))
with jax.debug_infs(True):
    e_inf = show_error(lambda: jax.jit(lambda x: 1.0 / x)(jnp.array(0.0)))
print("outside the block:", jax.config.jax_debug_nans, jax.config.jax_debug_infs)
assert isinstance(e_nan, FloatingPointError) and "nan" in str(e_nan)
assert isinstance(e_inf, FloatingPointError) and "inf" in str(e_inf)
assert jax.config.jax_debug_nans is False and jax.config.jax_debug_infs is False

NAN_PROBE = """
import sys, jax, jax.numpy as jnp
mode = sys.argv[1]
if mode == "with":
    jax.debug_nans(True).__enter__()               # equivalent to putting everything below in a with block
else:
    jax.config.update("jax_debug_nans", True)
f_old = jax.jit(lambda x: jnp.log(x))
f_old(jnp.array(2.0)).block_until_ready()          # run it normally once first
jnp.log(jnp.array(2.0))                            # also run the eager log once first
if mode in ("vmap", "clear", "with"):
    jax.vmap(lambda x: x * 2)(jnp.ones(3))         # one unrelated vmap
if mode == "clear":
    jax.clear_caches()                             # the workaround
cases = {
    "jit fn already run": lambda: f_old(jnp.array(-1.0)),
    "eager log already run": lambda: jnp.log(jnp.array(-1.0)),
    "eager sqrt never run": lambda: jnp.sqrt(jnp.array(-1.0)),
    "newly created jit fn": lambda: jax.jit(lambda x: jnp.arccos(x))(jnp.array(3.0)),
}
for name, fn in cases.items():
    try:
        fn(); print("RESULT", name, "missed")
    except FloatingPointError:
        print("RESULT", name, "caught")
"""
nan_table = {}
for mode in ["none", "vmap", "clear", "with"]:
    for line in result_lines(run_py(NAN_PROBE, mode)):
        name, verdict = line.rsplit(" ", 1)
        nan_table.setdefault(name, {})[mode] = verdict
print(f"{'Case':24s} {'no vmap':10s} {'vmap used':10s} {'vmap+clear':10s} {'with + vmap':10s}")
for name, row in nan_table.items():
    print(f"{name:22s} {row['none']:12s} {row['vmap']:12s} {row['clear']:12s} {row['with']}")

assert all(row["none"] == row["clear"] == row["with"] == "caught" for row in nan_table.values())
assert nan_table["jit fn already run"]["vmap"] == "missed" and nan_table["eager log already run"]["vmap"] == "missed"
assert nan_table["eager sqrt never run"]["vmap"] == "caught" and nan_table["newly created jit fn"]["vmap"] == "caught"

TRIGGER_PROBE = """
import sys, jax, jax.numpy as jnp
from jax.experimental import checkify
jax.config.update("jax_debug_nans", True)
f = jax.jit(lambda x: jnp.log(x))
f(jnp.array(2.0)).block_until_ready()
m = sys.argv[1]
if m == "grad": jax.grad(lambda x: x * 2)(1.0)
if m == "scan": jax.lax.scan(lambda c, x: (c + x, x), 0.0, jnp.ones(3))
if m == "make_jaxpr": jax.make_jaxpr(lambda x: x * 2)(1.0)
if m == "checkify": checkify.checkify(lambda x: x * 2)(1.0)
if m == "config.update": jax.config.update("jax_debug_nans", True)
if m == "make_jaxpr(vmap)": jax.make_jaxpr(jax.vmap(lambda x: x * 2))(jnp.ones(3))
if m == "jit(vmap)": jax.jit(jax.vmap(lambda x: x * 2))(jnp.ones(3))
try:
    f(jnp.array(-1.0)); print("RESULT missed")
except FloatingPointError:
    print("RESULT caught")
"""
trig = {m: result_lines(run_py(TRIGGER_PROBE, m))[0]
        for m in ["grad", "scan", "make_jaxpr", "checkify", "config.update", "make_jaxpr(vmap)", "jit(vmap)"]}
for m, v in trig.items():
    print(f"  after {m:18s} → NaN in the old jit function: {v}")
assert [m for m, v in trig.items() if v == "missed"] == ["make_jaxpr(vmap)", "jit(vmap)"]

def reg_loss(p, x):
    return jnp.sum((x @ p["w"] + p["b"]) ** 2) + 1e-3 * jnp.linalg.norm(p["b"])   # b is initialized to 0

p_reg = {"w": jnp.ones((3, 2)), "b": jnp.zeros(2)}
x_reg = jnp.ones((4, 3))
with jax.debug_nans(True):
    jax.clear_caches()
    err = show_error(lambda: jax.grad(reg_loss)(p_reg, x_reg))
user_frames = [f for f in traceback.extract_tb(err.__traceback__) if "site-packages" not in f.filename]
print("User code frames in the traceback:", [(f.name, f.line) for f in user_frames])

# Check leaf by leaf: which parameter's gradient contains NaN
grads = jax.grad(reg_loss)(p_reg, x_reg)
bad_leaves = [jax.tree_util.keystr(path) for path, g in jax.tree_util.tree_leaves_with_path(grads)
              if not bool(jnp.all(jnp.isfinite(g)))]
print("Leaves whose gradient contains NaN:", bad_leaves)
assert isinstance(err, FloatingPointError) and "mul" in str(err)
assert "reg_loss" not in [f.name for f in user_frames]        # no frame lands in reg_loss
assert bad_leaves == ["['b']"]

def risky(x, i):
    checkify.check(jnp.all(x > 0), "x must be positive, min = {m}", m=jnp.min(x))
    y = jnp.log(x)                  # x ≤ 0 → NaN
    z = x[i]                        # i out of bounds → JAX clamps the index by default, no error
    w = 7 // i                      # i == 0 → integer division by zero
    return y.sum() + z + w

x_ok = jnp.array([1.0, 2.0, 3.0])
inputs = {"normal": (x_ok, 1), "x negative": (jnp.array([-1.0, 2.0, 3.0]), 1), "index OOB": (x_ok, 5), "div by 0": (x_ok, 0)}
sets = {"user": checkify.user_checks, "nan": checkify.nan_checks, "index": checkify.index_checks,
        "div": checkify.div_checks, "all": checkify.all_checks}
table = {}
for sname, errs in sets.items():
    checked = jax.jit(checkify.checkify(risky, errors=errs))
    for iname, args in inputs.items():
        err, _ = checked(*args)
        table[(sname, iname)] = err.get()
print(f"{'Check set':8s}" + "".join(f"{n:>10s}" for n in inputs))
for sname in sets:
    print(f"{sname:10s}" + "".join(f"{'✔ caught' if table[(sname, n)] else '—':>10s}" for n in inputs))
print("\nSample messages:")
for key in [("user", "x negative"), ("nan", "x negative"), ("index", "index OOB"), ("div", "div by 0")]:
    print(f"  {key}: {table[key].strip()[:90]}")

assert all(table[(s, "normal")] is None for s in sets)
assert "x must be positive, min = -1.0" in table[("user", "x negative")]
assert "nan generated by primitive: log" in table[("nan", "x negative")]
assert "out-of-bounds" in table[("index", "index OOB")] and table[("user", "index OOB")] is None
assert "division by zero" in table[("div", "div by 0")]
assert all(table[("all", n)] is not None for n in inputs if n != "normal")
clamped = float(jax.jit(lambda x, i: x[i])(x_ok, 5))
print("\nActual result of x[5] without checks:", clamped, "(the index is clamped to 2, a silent error)")
assert clamped == 3.0

# ❌ Wrong: using checkify.check directly inside jit without wrapping with checkify.checkify
def check_only(x):
    checkify.check(x > 0, "x must be positive")
    return jnp.log(x)

try:
    jax.jit(check_only)(1.0)
    raised = None
except Exception as e:
    raised = e
    print(f"Actual error: {type(e).__name__}: {str(e).splitlines()[0][:120]}")
assert isinstance(raised, ValueError) and "not functionalized" in str(raised)

# ✅ Correct: checkify first, then jit; once you have err, inspect it with err.get() or raise it with err.throw()
checked_only = jax.jit(checkify.checkify(check_only))
err_ok, _ = checked_only(1.0)
err_bad, _ = checked_only(-1.0)
print("Good input:", err_ok.get(), "| bad input:", err_bad.get())
try:
    err_bad.throw()
except Exception as e:
    thrown = e
    print(f"err.throw() → {type(e).__name__} (inheritance chain {[c.__name__ for c in type(e).__mro__[:3]]})")
assert err_ok.get() is None and "x must be positive" in err_bad.get()
assert isinstance(thrown, checkify.JaxRuntimeError) and isinstance(thrown, ValueError)

def pos_log(x):
    checkify.check(x > 0, "x must be positive, got {v}", v=x)
    return jnp.log(x)

xs_mixed = jnp.array([1.0, -2.0, 3.0, -4.0])
r = {}
r["jit(checkify)"] = jax.jit(checkify.checkify(pos_log))(jnp.float32(-1.0))[0].get()
r["checkify(jit)"] = checkify.checkify(jax.jit(pos_log))(jnp.float32(-1.0))[0].get()
r["vmap(checkify)"] = jax.vmap(checkify.checkify(pos_log))(xs_mixed)[0].get()
r["checkify(vmap)"] = checkify.checkify(jax.vmap(pos_log))(xs_mixed)[0].get()
bad_where = lambda x: jnp.sum(jnp.where(x > 0, jnp.sqrt(x), 0.0))          # single where: the backward pass produces NaN
r["checkify(grad)"] = checkify.checkify(jax.grad(bad_where), errors=checkify.nan_checks)(jnp.array([0.0, 1.0]))[0].get()
g_inside = jax.grad(lambda x: checkify.checkify(pos_log)(x)[1])(jnp.float32(-1.0))
def scan_body(c, x):
    checkify.check(x < 10, "value at step {i} is too large: {x}", i=c[1], x=x)
    return (c[0] + x, c[1] + 1), None
r["scan"] = jax.jit(checkify.checkify(lambda xs: lax.scan(scan_body, (0.0, 0), xs)[0]))(jnp.array([1.0, 20.0, 30.0]))[0].get()
r["float_checks and inf"] = checkify.checkify(jnp.exp, errors=checkify.float_checks)(jnp.float32(1000.0))[0].get()
@jax.jit
def dbg_fn(x):
    checkify.debug_check(x > 0, "debug_check failed: {x}", x=x)
    return x * 2
dbg_plain = float(dbg_fn(-1.0))
r["debug_check + checkify"] = checkify.checkify(dbg_fn)(-1.0)[0].get()
for k, v in r.items():
    print(f"  {k:22s} → {v.strip() if v else None}".replace("\n", " ⏎ "))
print("  gradient of grad(checkify(f)[1]):", g_inside, "| result of debug_check when not wrapped:", dbg_plain)

assert r["jit(checkify)"] == r["checkify(jit)"] and "-1.0" in r["jit(checkify)"]
assert "at mapped index 1" in r["vmap(checkify)"] and "at mapped index 3" in r["vmap(checkify)"]
assert "mapped index" not in r["checkify(vmap)"] and r["checkify(vmap)"].count("must be positive") == 1
assert "nan generated by primitive: mul" in r["checkify(grad)"]
assert float(g_inside) == -1.0                                  # it differentiates, but the check is thrown away
assert "value at step 1 is too large: 20.0" in r["scan"]
assert r["float_checks and inf"] is None
assert dbg_plain == -2.0 and "debug_check failed" in r["debug_check + checkify"]

e_py = show_error(lambda: checkify.checkify(pos_log)(-1.0))
e_arr, _ = checkify.checkify(pos_log)(jnp.asarray(-1.0))
assert isinstance(e_py, TypeError) and "PyTrees of arrays" in str(e_py)
assert "got -1.0" in e_arr.get()

# ❌ Wrong: the backward rule of the custom_vjp has the wrong sign (the derivative of softplus should be sigmoid(x))
@jax.custom_vjp
def softplus_bad(x):
    return jnp.logaddexp(x, 0.0)
softplus_bad.defvjp(lambda x: (softplus_bad(x), x), lambda x, g: (g * jax.nn.sigmoid(-x),))

x_cg = jnp.array([-1.0, 0.5, 2.0])
try:
    check_grads(softplus_bad, (x_cg,), order=1, modes=["rev"])
    cg_err = None
except AssertionError as e:
    cg_err = e
    print("check_grads found an error:", " | ".join(str(e).strip().splitlines()[:2]))
e_fwd = show_error(lambda: check_grads(softplus_bad, (x_cg,), order=1, modes=["fwd"]))
assert cg_err is not None and "VJP" in str(cg_err)
assert isinstance(e_fwd, TypeError) and "custom_vjp" in str(e_fwd)

# ✅ Correct: the backward rule uses sigmoid(x), and check_grads passes quietly (it prints nothing when it passes)
@jax.custom_vjp
def softplus_ok(x):
    return jnp.logaddexp(x, 0.0)
softplus_ok.defvjp(lambda x: (softplus_ok(x), x), lambda x, g: (g * jax.nn.sigmoid(x),))

check_grads(softplus_ok, (x_cg,), order=1, modes=["rev"])
assert np.allclose(jax.grad(lambda x: softplus_ok(x).sum())(x_cg), 1 / (1 + np.exp(-to64(x_cg))), atol=1e-6)
print("check_grads passed ✔")

E = jax.errors
for cls in [E.ConcretizationTypeError, E.TracerBoolConversionError, E.TracerIntegerConversionError,
            E.TracerArrayConversionError, E.NonConcreteBooleanIndexError, E.UnexpectedTracerError]:
    print(f"  {cls.__name__:30s} → {[c.__name__ for c in cls.__mro__[1:4]]}")
assert issubclass(E.TracerBoolConversionError, E.ConcretizationTypeError)
assert not issubclass(E.TracerIntegerConversionError, E.ConcretizationTypeError)
assert not issubclass(E.TracerArrayConversionError, E.ConcretizationTypeError)
assert all(issubclass(c, E.JAXTypeError) and issubclass(c, TypeError) for c in
           [E.ConcretizationTypeError, E.TracerIntegerConversionError, E.TracerArrayConversionError, E.UnexpectedTracerError])
assert issubclass(E.NonConcreteBooleanIndexError, IndexError)

def clip_positive(x, limit):
    if x > limit:                  # depends on the values of x and limit
        return limit
    return x

e = show_error(lambda: jax.jit(clip_positive)(jnp.float32(3.0), 1.0))
msg = " ".join(str(e).split())                       # the original message has line breaks; squash it into one line before matching
m = re.search(r"The error occurred while tracing the function (\S+) at \S+ for jit\..*?"
              r"depends on the values? of the arguments? (.+?)\.", msg)
print("Function at fault:", m.group(1), "| culprit argument(s):", m.group(2))
assert isinstance(e, E.TracerBoolConversionError)
assert m.group(1) == "clip_positive" and m.group(2) == "x and limit"   # both arguments take part in the comparison

# ❌ Wrong: a Python if that depends on a traced value
@jax.jit
def relu_if(x):
    if x > 0:
        return x
    return 0.0 * x

try:
    relu_if(jnp.float32(2.0))
    caught = None
except Exception as e:
    caught = e
    print(f"Actual error: {type(e).__name__}: {str(e).splitlines()[0]}")
assert isinstance(caught, E.TracerBoolConversionError)

# ✅ Correct: use jnp.where elementwise; use lax.cond when only one branch should run
@jax.jit
def relu_where(x):
    return jnp.where(x > 0, x, 0.0)

@jax.jit
def relu_cond(x):
    return lax.cond(x > 0, lambda x: x, lambda x: 0.0 * x, x)

xs = jnp.array([-2.0, 0.0, 3.0])
assert np.array_equal(relu_where(xs), [0.0, 0.0, 3.0])
assert [float(relu_cond(v)) for v in xs] == [0.0, 0.0, 3.0]
print("jnp.where:", relu_where(xs), "| lax.cond:", [float(relu_cond(v)) for v in xs])

# ❌ Wrong: using a traced value as the range count or as a Python list index
def repeat_add(x, n):
    for _ in range(n):
        x = x + 1.0
    return x

def pick(i):
    return [10.0, 20.0, 30.0][i]

for name, fn in [("range(n)", lambda: jax.jit(repeat_add)(1.0, 3)), ("list[i]", lambda: jax.jit(pick)(1))]:
    try:
        fn()
        caught = None
    except Exception as e:
        caught = e
        print(f"{name:9s} actual error: {type(e).__name__}: {str(e).splitlines()[0]}")
    assert isinstance(caught, E.TracerIntegerConversionError)

# ✅ Correct: the count determines the computation structure → static argument; or use fori_loop; convert the list to an array before indexing
repeat_static = jax.jit(repeat_add, static_argnames="n")
repeat_fori = jax.jit(lambda x, n: lax.fori_loop(0, n, lambda _, v: v + 1.0, x))
pick_arr = jax.jit(lambda i: jnp.asarray([10.0, 20.0, 30.0])[i])
assert float(repeat_static(1.0, n=3)) == 4.0 and float(repeat_fori(1.0, 3)) == 4.0 and float(pick_arr(1)) == 20.0
print("static:", repeat_static(1.0, n=3), "| fori_loop:", repeat_fori(1.0, 3), "| array indexing:", pick_arr(1))