Chapter 7: Practical Topics

# Common imports and helpers (every later cell depends on this one)
import time, math, logging, warnings, functools, dataclasses
from dataclasses import dataclass, field
from functools import partial
from typing import NamedTuple, Any, Callable

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
import jax.monitoring as monitoring
import matplotlib.pyplot as plt

plt.rcParams["font.sans-serif"] = ["PingFang SC", "Heiti SC", "Arial Unicode MS", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False
# Chart colors (categorical colors in a fixed order)
C_BLUE, C_ORANGE, C_AQUA, C_VIOLET, C_RED = "#2a78d6", "#eb6834", "#1baf7a", "#6250d6", "#e34948"

print("JAX", jax.__version__, "| backend", jax.default_backend(), "| devices", jax.devices())


def bench(f, *args, warmup=2, runs=10):
    """Timing: warm up `warmup` times (including compilation), then run `runs` times and return the median in ms."""
    for _ in range(warmup):
        jax.block_until_ready(f(*args))
    ts = []
    for _ in range(runs):
        t0 = time.perf_counter()
        jax.block_until_ready(f(*args))
        ts.append(time.perf_counter() - t0)
    return float(np.median(ts)) * 1e3


def assert_faster(fast, slow, factor=1.0, tries=3):
    """Timing assertion: time of fast() × factor < time of slow() (both arguments are no-arg functions that return a time).
    Timings depend on machine load (e.g. several notebooks running in parallel on CI), so if it fails, measure again, up to `tries` times."""
    for _ in range(tries):
        tf, ts = fast(), slow()
        if tf * factor < ts:
            return tf, ts
    raise AssertionError(f"Still not satisfied after {tries} tries: {tf:.3g} × {factor} ≥ {ts:.3g}")

def show_error(fn):
    """Run fn; if it raises, print only the exception type and the first line of the message (for demonstrating wrong code)."""
    try:
        out = fn()
        print("(no error) result:", out)
    except Exception as e:
        first = str(e).strip().splitlines()[0] if str(e).strip() else ""
        print(f"❌ {type(e).__name__}: {first[:160]}")


_COMPILE_EVENT = "/jax/core/compile/backend_compile_duration"


class CompileCounter:
    """Count the XLA backend compilations that happen inside a with block.

    name=None counts all compilations; name="f" counts only compilations of jit(f).
    How it works: every time JAX finishes a backend compilation, it emits a
    '/jax/core/compile/backend_compile_duration' event through jax.monitoring, with fun_name attached.
    """
    def __init__(self, name=None):
        self.name = name
        self.events = []                      # [(fun_name, seconds)]
        self._listener = self._on_event       # keep the same bound-method object so we can find it when unregistering

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

    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 wpad(s, width):
    """Left-align by display width (a CJK character takes 2 columns), for printing aligned tables."""
    import unicodedata
    w = sum(2 if unicodedata.east_asian_width(ch) in "WF" else 1 for ch in str(s))
    return str(s) + " " * max(0, width - w)


def to64(x):
    """Convert to NumPy float64, for cross-checking against high-precision reference values."""
    return np.asarray(x, dtype=np.float64)

# Self-test: the counter sees compilations and doesn't count cache hits
_x = jnp.ones(3)
def _probe(x):
    return x * 2
with CompileCounter("_probe") as _c:
    jax.jit(_probe)(_x); jax.jit(_probe)(_x)
assert _c.count == 1, _c.events
print("CompileCounter self-test passed: two calls, only 1 compilation")

def mlp_loss(w1, w2, x):
    # Forward pass of a small MLP + mean square; weights and data are passed in as arguments
    h = jnp.tanh(x @ w1)
    return jnp.mean((h @ w2) ** 2)

k1, k2, k3 = jax.random.split(jax.random.key(0), 3)
W1 = jax.random.normal(k1, (64, 256)) / 8.0
W2 = jax.random.normal(k2, (256, 10)) / 16.0
X_all = jax.random.normal(k3, (600, 64))
grad_fn = jax.jit(jax.grad(mlp_loss, argnums=(0, 1)))

sizes = [100 + 7 * i for i in range(20)]            # 20 different batch sizes
batches = [X_all[:n] for n in sizes]                # slice ahead of time so slicing isn't timed
same = [X_all[:100]] * 20

def run(xs):
    t0 = time.perf_counter()
    for x in xs:
        jax.block_until_ready(grad_fn(W1, W2, x))
    return (time.perf_counter() - t0) * 1e3

with CompileCounter("mlp_loss") as c_var:
    t_var = run(batches)
grad_fn(W1, W2, same[0])                             # warm up: make sure only execution is measured below
with CompileCounter("mlp_loss") as c_same:
    t_same = run(same)
print(f"20 batch sizes:      {t_var:7.1f} ms, {c_var.count} compilations")
print(f"same batch size:     {t_same:7.1f} ms, {c_same.count} compilations")
print(f"About {np.mean([d for _, d in c_var.events]) * 1e3:.1f} ms per compilation, "
      f"about {t_same / 20:.2f} ms per execution")
assert c_var.count == 20 and c_same.count == 0

x3, x4 = jnp.ones(3), jnp.ones(4)
x3_i = jnp.ones(3, jnp.int32)

def count_compiles(name, calls):
    """Run each no-arg function in calls in order, and return how many times jit(name) was compiled."""
    with CompileCounter(name) as c:
        for call in calls:
            jax.block_until_ready(call())
    return c.count

rows = []   # (trigger, how it's called, measured compile count, expected count)

# ① Shape
def t_shape(x): return x * 2
f = jax.jit(t_shape)
rows.append(("shape", "f32[3] → f32[4]", count_compiles("t_shape", [lambda: f(x3), lambda: f(x4)]), 2))

# ② dtype
def t_dtype(x): return x * 2
f = jax.jit(t_dtype)
rows.append(("dtype", "f32[3] → i32[3]", count_compiles("t_dtype", [lambda: f(x3), lambda: f(x3_i)]), 2))

# ③ Weak type: Python scalars are "weakly typed", arrays are "strongly typed"
def t_weak(x, lr): return x * lr
f = jax.jit(t_weak)
rows.append(("weak type", "lr=0.1 → lr=jnp.float32(0.1)",
             count_compiles("t_weak", [lambda: f(x3, 0.1), lambda: f(x3, jnp.float32(0.1))]), 2))

# ④ Value of a static argument
def t_static(x, n): return x ** n
f = jax.jit(t_static, static_argnames="n")
rows.append(("static arg value", "n=2 → n=3", count_compiles("t_static", [lambda: f(x3, n=2), lambda: f(x3, n=3)]), 2))

# ⑤ Pytree structure: one more key / list replaced by tuple / leaf replaced by None
def t_tree(s): return jax.tree.map(lambda a: a + 1, s)
f = jax.jit(t_tree)
rows.append(("pytree structure", "{'a'} → {'a','b'} → [x] → (x,) → {'a', 'b':None}",
             count_compiles("t_tree", [lambda: f({"a": x3}), lambda: f({"a": x3, "b": x3}),
                                       lambda: f([x3]), lambda: f((x3,)), lambda: f({"a": x3, "b": None})]), 5))

# ⑥ New jit object: every call creates a new inner function (closure)
def make_and_call(x):
    def t_inner(y): return y + 1          # every call to make_and_call creates a new function object
    return jax.jit(t_inner)(x)
rows.append(("new jit obj (closure)", "called 3 times in a loop",
             count_compiles("t_inner", [lambda: make_and_call(x3)] * 3), 3))

# ⑦ New jit object: jit a bound method (every access to obj.method creates a new bound-method object)
class Model:
    def t_method(self, x): return x * 3
    def predict(self, x): return jax.jit(self.t_method)(x)
m = Model()
rows.append(("new jit obj (bound meth)", "m.predict called 3 times",
             count_compiles("t_method", [lambda: m.predict(x3)] * 3), 3))

# ⑧ Static arguments compared by id: a custom class without __eq__/__hash__ makes every new instance a new key
class Cfg:
    def __init__(self, n): self.n = n
def t_cfg(x, cfg): return x * cfg.n
f = jax.jit(t_cfg, static_argnums=1)
rows.append(("static arg hashed by id", "Cfg(2) created 3 times", count_compiles("t_cfg", [lambda: f(x3, Cfg(2))] * 3), 3))

# ⑨ How you call it: omitting a default argument vs. passing the same value explicitly
def t_default(x, scale=2.0): return x * scale
f = jax.jit(t_default)
rows.append(("calling convention", "f(x) → f(x, 2.0) → f(x, scale=2.0)",
             count_compiles("t_default", [lambda: f(x3), lambda: f(x3, 2.0), lambda: f(x3, scale=2.0)]), 3))

# -- Control group: these don't trigger recompilation --
def t_value(x): return x * 2
f = jax.jit(t_value)
rows.append(("(control) values only", "ones(3) → ones(3)*5", count_compiles("t_value", [lambda: f(x3), lambda: f(x3 * 5)]), 1))

def t_shared(x): return x * 2
rows.append(("(control) jax.jit(f) x3", "same module-level function wrapped 3 times",
             count_compiles("t_shared", [lambda: jax.jit(t_shared)(x3)] * 3), 1))

print(wpad("Trigger", 26) + wpad("Calls", 50) + " Got   Exp")
print("-" * 80)
for name, how, got, exp in rows:
    print(wpad(name, 26) + wpad(how, 50) + f"{got:>4}{exp:>6}")
for name, how, got, exp in rows:
    assert got == exp, (name, got, exp)
print("\nAll match the expectations ✔")

class capture_jax_logs:
    """Temporarily take over the 'jax' logger and collect log messages into self.messages (no longer printed 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

# ① jax_log_compiles: prints one "Compiling jit(...)" line for every real compilation
def diag_step(params, x):
    return jax.tree.map(lambda p: p * x.mean(), params)

f = jax.jit(diag_step)
params_a = {"w": jnp.ones((2, 2)), "b": jnp.zeros(2)}
xs = [jnp.ones(5), jnp.ones(5) * 2, jnp.ones(6)]           # build the inputs first so jnp.ones's own compile logs don't get mixed in

jax.config.update("jax_log_compiles", True)
try:
    with capture_jax_logs() as logs:
        for x in xs:
            f(params_a, x)
finally:
    jax.config.update("jax_log_compiles", False)            # finally makes sure it gets turned off
compiling = [m for m in logs.messages if m.startswith("Compiling")]
for m in compiling:
    print(m[:150], "...")
assert len(compiling) == 2          # the 2nd call only has different values, so it hits the cache

# ② jax_explain_cache_misses: tells you directly "which argument changed, and where"
def diag_explain(state, lr):
    return jax.tree.map(lambda p: p - lr * p, state)

f = jax.jit(diag_explain)
s1 = {"w": jnp.ones(3)}
s2 = {"w": jnp.ones(3), "m": jnp.zeros(3)}                  # the structure changed
f(s1, 0.1)                                                  # populate the cache first
jax.config.update("jax_explain_cache_misses", True)
try:
    with capture_jax_logs() as logs:
        f(s1, jnp.float32(0.1))        # weak type → strong type
        f(s2, 0.1)                     # pytree structure change
finally:
    jax.config.update("jax_explain_cache_misses", False)
misses = [m for m in logs.messages if "CACHE MISS" in m]
for m in misses:
    # drop the file-path lines and keep only the "reason" part
    lines = [l for l in m.splitlines() if not l.strip().startswith(("for ", "TRACING"))]
    print("\n".join(lines[:6]), "\n")
assert len(misses) == 2

trace_count = 0
def diag_counter(x):
    global trace_count
    trace_count += 1               # runs only during tracing
    return jnp.sin(x) * 2

f = jax.jit(diag_counter)
with CompileCounter("diag_counter") as c:
    jax.eval_shape(f, x3)          # trace only, compute the output shape/dtype
    f.lower(x4)                    # trace + lower to StableHLO, no compile
    f(x3); f(x3)                   # 1st: trace + compile (reusing the trace from eval_shape); 2nd: cache hit
print(f"trace count = {trace_count}, backend compile count = {c.count}")
assert c.count == 1 and trace_count >= 2

def deep_unrolled(ws, x):
    for w in ws:                   # the Python loop gets unrolled into 30 copies during tracing
        x = jnp.tanh(x @ w)
    return x.sum()

ws = [jax.random.normal(k, (64, 64)) / 8.0 for k in jax.random.split(jax.random.key(1), 30)]
xb = jnp.ones((16, 64))
jf = jax.jit(deep_unrolled)

t0 = time.perf_counter(); traced = jf.trace(ws, xb); t1 = time.perf_counter()
lowered = traced.lower();                                t2 = time.perf_counter()
compiled = lowered.compile();                            t3 = time.perf_counter()
print(f"trace {1e3*(t1-t0):6.1f} ms | lower {1e3*(t2-t1):6.1f} ms | compile {1e3*(t3-t2):6.1f} ms")
print("jaxpr equations:", len(traced.jaxpr.eqns), "| StableHLO lines:", len(lowered.as_text().splitlines()))
print("Types:", type(traced).__name__, type(lowered).__name__, type(compiled).__name__,
      "| all in jax.stages:", all(hasattr(jax.stages, n) for n in ("Traced", "Lowered", "Compiled")))

ref = to64(xb)                                            # NumPy float64 reference value
for w in ws:
    ref = np.tanh(ref @ to64(w))
assert np.allclose(compiled(ws, xb), ref.sum(), rtol=1e-4, atol=1e-4)
show_error(lambda: compiled(ws, jnp.ones((17, 64))))       # different shape → TypeError

def masked_mse(w, X, y, mask):
    r = X @ w - y
    return jnp.sum(mask * r ** 2) / jnp.sum(mask)          # numerator and denominator count only valid rows

def plain_mse(w, X, y):
    return jnp.mean((X @ w - y) ** 2)

BUCKETS = (64, 128, 256, 512)

def pad_to_bucket(X, y):
    """Pad in NumPy (triggers no JAX compilation) and return (Xp, yp, mask)."""
    n = X.shape[0]
    b = next(b for b in BUCKETS if b >= n)
    Xp = np.zeros((b, X.shape[1]), np.float32); Xp[:n] = X
    yp = np.zeros(b, np.float32); yp[:n] = y
    mask = np.zeros(b, np.float32); mask[:n] = 1.0
    return Xp, yp, mask

kx, ky, kw, kn = jax.random.split(jax.random.key(2), 4)
Xbig = np.asarray(jax.random.normal(kx, (512, 8)))
ybig = np.asarray(jax.random.normal(ky, (512,)))
w0 = jax.random.normal(kw, (8,))
lengths = [int(n) for n in jax.random.randint(kn, (30,), 20, 500)]   # 30 random lengths

vg_plain = jax.jit(jax.value_and_grad(plain_mse))
vg_masked = jax.jit(jax.value_and_grad(masked_mse))
wrong_mean = jax.jit(lambda w, X, y: jnp.mean((X @ w - y) ** 2))      # ❌ padded but not masked

with CompileCounter("plain_mse") as c_plain:
    for n in lengths:
        vg_plain(w0, Xbig[:n], ybig[:n])
max_err_masked, max_err_wrong = 0.0, 0.0
with CompileCounter("masked_mse") as c_masked:
    for n in lengths:
        Xp, yp, mask = pad_to_bucket(Xbig[:n], ybig[:n])
        loss, g = vg_masked(w0, Xp, yp, mask)
        # NumPy float64 reference values
        r = Xbig[:n].astype(np.float64) @ to64(w0) - ybig[:n]
        ref_loss, ref_g = np.mean(r ** 2), 2 * Xbig[:n].T.astype(np.float64) @ r / n
        max_err_masked = max(max_err_masked, abs(float(loss) - ref_loss), np.abs(to64(g) - ref_g).max())
        max_err_wrong = max(max_err_wrong, abs(float(wrong_mean(w0, Xp, yp)) - ref_loss))

print(f"No buckets: {len(set(lengths))} lengths → {c_plain.count} compilations")
print(f"Buckets:    {len(BUCKETS)} buckets    → {c_masked.count} compilations")
print(f"Max error of buckets + mask vs. NumPy float64:  {max_err_masked:.2e}")
print(f"❌ Max loss error when padding without a mask:   {max_err_wrong:.3f}")
assert c_plain.count == len(set(lengths)) and c_masked.count <= len(BUCKETS)
assert max_err_masked < 1e-4 and max_err_wrong > 0.1

# The cost of padding: how much computation is wasted on average
waste = np.mean([1 - n / next(b for b in BUCKETS if b >= n) for n in lengths])
print(f"Power-of-2 buckets waste {waste:.0%} of the computation on average; finer buckets waste less but compile more often")

@dataclass(frozen=True)          # frozen=True → automatically generates value-based __eq__ and __hash__
class CfgOK:
    n: int

def t_cfg_ok(x, cfg): return x * cfg.n
f = jax.jit(t_cfg_ok, static_argnums=1)
n_ok = count_compiles("t_cfg_ok", [lambda: f(x3, CfgOK(2))] * 3)

@jax.tree_util.register_static   # register as a static pytree: just pass it in as a regular argument
@dataclass(frozen=True)
class Activation:
    name: str

def t_act(x, act):
    return jnp.tanh(x) if act.name == "tanh" else jax.nn.relu(x)   # Python if: act is static
f_act = jax.jit(t_act)           # note: no static_argnums
n_act = count_compiles("t_act", [lambda: f_act(x3, Activation("tanh"))] * 3 + [lambda: f_act(x3, Activation("relu"))])

# Learning rate: ❌ static (one compilation per new value) vs. ✅ dynamic (compiled once)
def t_lr_static(p, lr): return p - lr * p
def t_lr_dynamic(p, lr): return p - lr * p
f_s = jax.jit(t_lr_static, static_argnames="lr")
f_d = jax.jit(t_lr_dynamic)
lrs = [1e-3 * (1 + i) for i in range(8)]                      # simulate 8 steps of a learning-rate schedule
n_s = count_compiles("t_lr_static", [lambda lr=lr: f_s(x3, lr=lr) for lr in lrs])
n_d = count_compiles("t_lr_dynamic", [lambda lr=lr: f_d(x3, jnp.float32(lr)) for lr in lrs])

print(f"frozen dataclass as a static argument, 3 equal instances → {n_ok} compilation(s)")
print(f"register_static: tanh × 3 + relu × 1                 → {n_act} compilation(s)")
print(f"8-step learning-rate schedule: static {n_s} | dynamic {n_d} compilations")
assert (n_ok, n_act, n_s, n_d) == (1, 2, 8, 1)

def _forward(w, x):                         # pure function: weights passed explicitly
    return jnp.tanh(x @ w)
forward_jit = jax.jit(_forward)             # ✅ created once at module level

class ModelOK:
    def __init__(self, w):
        self.w = w
    def predict(self, x):
        return forward_jit(self.w, x)       # the method only calls it, it doesn't create a new jit

mo = ModelOK(jnp.ones((3, 2)))
n_fwd = count_compiles("_forward", [lambda: mo.predict(x3)] * 5)

def t_lr_mixed(p, lr): return p - lr * p
f_mix = jax.jit(t_lr_mixed)
schedule = [0.1, jnp.float32(0.05), 0.02]                     # ❌ mixing Python floats and arrays
n_mix = count_compiles("t_lr_mixed", [lambda lr=lr: f_mix(x3, lr) for lr in schedule])
def t_lr_cast(p, lr): return p - lr * p
f_cast = jax.jit(t_lr_cast)
n_cast = count_compiles("t_lr_cast", [lambda lr=lr: f_cast(x3, jnp.asarray(lr, jnp.float32)) for lr in schedule])  # ✅
def t_lr_py(p, lr): return p - lr * p
f_py = jax.jit(t_lr_py)
n_py = count_compiles("t_lr_py", [lambda lr=lr: f_py(x3, lr) for lr in [0.1, 0.05, 0.02]])  # all Python floats: no recompilation even though the value changes

print(f"Module-level jit, predict called 5 times → {n_fwd} compilation(s)")
print(f"Learning rate mixing Python float / array → {n_mix} compilations; all converted to float32 arrays → {n_cast}; all Python floats → {n_py} compilation(s)")
assert (n_fwd, n_mix, n_cast, n_py) == (1, 2, 1, 1)

def serve_loss(w, X, y, mask):
    return masked_mse(w, X, y, mask)

t0 = time.perf_counter()
compiled_by_bucket = {
    b: jax.jit(serve_loss).lower(
        jax.ShapeDtypeStruct((8,), jnp.float32), jax.ShapeDtypeStruct((b, 8), jnp.float32),
        jax.ShapeDtypeStruct((b,), jnp.float32), jax.ShapeDtypeStruct((b,), jnp.float32)).compile()
    for b in BUCKETS}
print(f"Precompiling {len(BUCKETS)} buckets at startup: {(time.perf_counter() - t0) * 1e3:.0f} ms")

w0_np = np.asarray(w0)
with CompileCounter() as c_serve:                 # name=None: count compilations of all functions
    outs = []
    for n in lengths:
        Xp, yp, mask = pad_to_bucket(Xbig[:n], ybig[:n])      # padding done in NumPy
        outs.append(compiled_by_bucket[Xp.shape[0]](w0_np, Xp, yp, mask))
    jax.block_until_ready(outs)
print(f"Serving {len(lengths)} requests triggered {c_serve.count} compilations")
assert c_serve.count == 0

class CompileBudgetExceeded(RuntimeError):
    pass


class CompileBudget(CompileCounter):
    """Compile budget: raise (or warn) when there are more than max_compiles compilations."""
    def __init__(self, max_compiles, name=None, on_exceed="raise"):
        super().__init__(name)
        self.max_compiles = max_compiles
        self.on_exceed = on_exceed

    def check(self):
        if self.count > self.max_compiles:
            detail = ", ".join(f"{fn}({d * 1e3:.0f}ms)" for fn, d in self.events)
            msg = (f"Compile budget exceeded: {self.name or 'all functions'} allows {self.max_compiles}, "
                   f"got {self.count} → {detail}")
            if self.on_exceed == "raise":
                raise CompileBudgetExceeded(msg)
            warnings.warn(msg, RuntimeWarning, stacklevel=3)

    def __exit__(self, exc_type, exc, tb):
        super().__exit__(exc_type, exc, tb)
        if exc_type is None:          # if the with block itself raised, don't pile a budget error on top
            self.check()
        return False


def compile_budget(max_compiles, on_exceed="raise"):
    """Decorate a jit function: accumulate its compile count and check the budget after every call."""
    def deco(jitted):
        budget = CompileBudget(max_compiles, jitted.__name__, on_exceed)

        @functools.wraps(jitted)
        def wrapper(*args, **kwargs):
            monitoring.register_event_duration_secs_listener(budget._listener)
            try:
                out = jitted(*args, **kwargs)
            finally:
                monitoring.unregister_event_duration_listener(budget._listener)
            budget.check()
            return out

        wrapper.budget = budget
        return wrapper
    return deco

# Usage 1: decorator. Allow 2 compilations (e.g. one batch shape each for training and evaluation)
@compile_budget(2)
@jax.jit
def budget_step(w, x):
    return w - 0.1 * jnp.mean(x) * w

w_ = jnp.ones(4)
for n in [32, 32, 64, 64, 32]:                 # only 2 shapes: within budget
    w_ = budget_step(w_, jnp.ones(n))
print("5 calls, 2 shapes:", budget_step.budget.count, "compilations, within budget")
assert budget_step.budget.count == 2

try:
    budget_step(w_, jnp.ones(48))               # a 3rd shape → over budget
    raise AssertionError("should have raised")
except CompileBudgetExceeded as e:
    print("❌ Caught:", str(e)[:120])

# Usage 2: context manager + warn mode
def budget_loop(p, lr): return p - lr * p
f = jax.jit(budget_loop)
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    with CompileBudget(1, "budget_loop", on_exceed="warn") as b:
        for lr in [0.1, 0.05, jnp.float32(0.01)]:   # one strongly typed value mixed in → one extra compilation
            f(x3, lr)
print("⚠️ Warning:", str(caught[0].message)[:110])
assert b.count == 2 and len(caught) == 1

# Usage 3: ✅ as an assertion: once fixed, it stays within budget and passes quietly
with CompileBudget(1, "budget_loop"):
    for lr in [0.1, 0.05, 0.01]:
        f(x3, jnp.asarray(lr, jnp.float32))         # strongly typed f32[] was already compiled above → 0
print("After the fix: within budget ✔")

# One NaN example wipes out the whole set of parameters in a single step
def per_sample_loss(w, x):
    return jnp.sqrt(jnp.sum((w * x) ** 2))      # ‖w ⊙ x‖: when x is all zeros the gradient is 0/0

w_nan = jnp.ones(4)
X_nan = jnp.stack([jnp.ones(4), 2 * jnp.ones(4), jnp.zeros(4)])     # the 3rd example is all zeros
g = jax.grad(lambda w: jnp.mean(jax.vmap(per_sample_loss, (None, 0))(w, X_nan)))(w_nan)
w_after = w_nan - 0.1 * g
print("float32 exp limits: exp(88.7) =", jnp.exp(88.7), "| exp(89) =", jnp.exp(89.0),
      "| exp(-87) =", jnp.exp(-87.0), "| exp(-88) =", jnp.exp(-88.0), "| NumPy exp(-100) =", np.exp(np.float32(-100.0)))
assert float(jnp.exp(-88.0)) == 0.0 and float(jnp.exp(-87.0)) > 0 and np.exp(np.float32(-100.0)) > 0   # JAX flushes to zero, NumPy keeps denormals
print("Batch gradient:", g)
print("Parameters after one update:", w_after)
assert bool(jnp.all(jnp.isnan(w_after)))

def double_where(fn, x, ok, safe_value, fill):
    """Return fill where ok is False, while keeping the gradient there finite."""
    x_safe = jnp.where(ok, x, safe_value)
    return jnp.where(ok, fn(x_safe), fill)

# ---- 4 dangerous functions: ❌ single-where version vs. ✅ double-where version ----
def norm_bad(v):                                   # vector norm, 0/0 at the zero vector
    return jnp.linalg.norm(v)
def norm_good(v):
    sq = jnp.sum(v ** 2)
    return double_where(jnp.sqrt, sq, sq > 0, 1.0, 0.0)

def sinc_bad(x):                                   # sin(x)/x, filled with 1 at x=0
    return jnp.where(x != 0, jnp.sin(x) / x, 1.0)
def sinc_good(x):
    return double_where(lambda t: jnp.sin(t) / t, x, x != 0, 1.0, 1.0)

def entropy_bad(p):                                # -Σ p log p, defined as 0 at p=0
    return -jnp.sum(jnp.where(p > 0, p * jnp.log(p), 0.0))
def entropy_good(p):
    return -jnp.sum(double_where(lambda t: t * jnp.log(t), p, p > 0, 1.0, 0.0))

def logit_bad(p):                                  # log(p/(1-p)), clipped to ±20 for p∈{0,1}
    return jnp.where((p > 0) & (p < 1), jnp.log(p / (1 - p)), 20.0 * jnp.sign(p - 0.5))
def logit_good(p):
    ok = (p > 0) & (p < 1)
    return double_where(lambda t: jnp.log(t / (1 - t)), p, ok, 0.5, 20.0 * jnp.sign(p - 0.5))

cases = [
    ("norm ‖v‖ @ v=0",        norm_bad,    norm_good,    jnp.zeros(3)),
    ("sinc @ x=0",            sinc_bad,    sinc_good,    jnp.array(0.0)),
    ("entropy @ p=[0,.5,.5]", entropy_bad, entropy_good, jnp.array([0.0, 0.5, 0.5])),
    ("logit @ p=1",           logit_bad,   logit_good,   jnp.array(1.0)),
]
print(wpad("Function @ danger pt", 24) + wpad("Forward (❌/✅)", 22) + "grad ❌  →  grad ✅")
for name, bad, good, x in cases:
    gb, gg = jax.grad(bad)(x), jax.grad(good)(x)
    print(wpad(name, 24) + wpad(f"{float(bad(x)):.3f} / {float(good(x)):.3f}", 22) + f"{gb}  →  {gg}")
    assert np.allclose(bad(x), good(x))                       # forward values are identical
    assert not bool(jnp.all(jnp.isfinite(gb)))                # ❌ gradient contains NaN
    assert bool(jnp.all(jnp.isfinite(gg)))                    # ✅ gradient is all finite

# At normal points, the gradient of the ✅ version must equal the analytic derivative (computed in NumPy float64)
v, x, p, q = np.array([3.0, 4.0, 12.0]), 0.7, np.array([0.2, 0.3, 0.5]), 0.3
assert np.allclose(jax.grad(norm_good)(jnp.array(v)), v / np.linalg.norm(v), atol=1e-6)
assert np.allclose(jax.grad(sinc_good)(x), (x * np.cos(x) - np.sin(x)) / x ** 2, atol=1e-6)
assert np.allclose(jax.grad(entropy_good)(jnp.array(p)), -(np.log(p) + 1), atol=1e-6)
assert np.allclose(jax.grad(logit_good)(q), 1 / (q * (1 - q)), atol=1e-4)
print("\nGradients at normal points match the analytic derivatives ✔")

def logsumexp_stable(x, axis=-1):
    m = jnp.max(x, axis=axis, keepdims=True)
    m = lax.stop_gradient(jnp.where(jnp.isfinite(m), m, 0.0))     # details 1 and 2
    return jnp.squeeze(m, axis) + jnp.log(jnp.sum(jnp.exp(x - m), axis=axis))

def log_softmax_stable(x, axis=-1):
    return x - jnp.expand_dims(logsumexp_stable(x, axis), axis)

def softmax_stable(x, axis=-1):
    m = lax.stop_gradient(jnp.max(x, axis=axis, keepdims=True))
    e = jnp.exp(x - m)
    return e / jnp.sum(e, axis=axis, keepdims=True)

# ❌ Naive versions
def logsumexp_naive(x): return jnp.log(jnp.sum(jnp.exp(x), axis=-1))
def softmax_naive(x):   return jnp.exp(x) / jnp.sum(jnp.exp(x), axis=-1, keepdims=True)
def log_softmax_naive(x): return jnp.log(softmax_naive(x))

def np_log_softmax(x):          # NumPy float64 reference
    x = np.asarray(x, np.float64)
    m = x.max(-1, keepdims=True)
    return x - (m + np.log(np.exp(x - m).sum(-1, keepdims=True)))

logits = jnp.array([[1000.0, 0.0, -1000.0],     # overflow
                    [-1000.0, -1001.0, -1002.0], # everything underflows
                    [0.0, -120.0, 3.0],          # one tiny probability
                    [1.0, 2.0, 3.0]])            # ordinary case
ref = np_log_softmax(logits)
print("Naive log(softmax):\n", log_softmax_naive(logits))
print("Stable log_softmax:\n", log_softmax_stable(logits))
print("Naive logsumexp:", logsumexp_naive(logits), "\nStable logsumexp:", logsumexp_stable(logits))

assert np.allclose(log_softmax_stable(logits), ref, atol=1e-4)
assert np.allclose(log_softmax_stable(logits), jax.nn.log_softmax(logits), atol=1e-5)
assert np.allclose(softmax_stable(logits), np.exp(ref), atol=1e-6)
assert np.allclose(logsumexp_stable(logits), jax.nn.logsumexp(logits, axis=-1), atol=1e-4)
assert not bool(jnp.all(jnp.isfinite(log_softmax_naive(logits))))
print("The stable versions match NumPy float64 and jax.nn ✔")

def ce_naive(z, y):
    return -jnp.mean(jnp.take_along_axis(log_softmax_naive(z), y[:, None], axis=1))

def ce_stable(z, y):
    return -jnp.mean(jnp.take_along_axis(log_softmax_stable(z), y[:, None], axis=1))

z = jnp.array([[0.0, -120.0, 3.0], [1.0, 2.0, 3.0]])
y = jnp.array([1, 2])                        # the correct-class probability of the 1st example is about e^-123
for name, fn in [("❌ log(softmax)", ce_naive), ("✅ log_softmax", ce_stable)]:
    loss, gz = jax.value_and_grad(fn)(z, y)
    print(f"{name}: loss = {float(loss):.4f}, gradient finite = {bool(jnp.all(jnp.isfinite(gz)))}")

loss, gz = jax.value_and_grad(ce_stable)(z, y)
ref_g = (np.exp(np_log_softmax(z)) - np.eye(3)[np.asarray(y)]) / len(y)       # (softmax - onehot) / B
ref_loss = -np.mean(np_log_softmax(z)[np.arange(2), np.asarray(y)])
assert np.allclose(gz, ref_g, atol=1e-6) and np.isclose(float(loss), ref_loss, rtol=1e-5)
assert not np.isfinite(float(ce_naive(z, y)))
print("Gradient = (softmax − onehot)/B, matches NumPy float64 ✔")

NEG = jnp.finfo(jnp.float32).min

def masked_ce(z, y, valid_cols, row_weight, fill):
    z = jnp.where(valid_cols, z, fill)                       # mask out invalid classes
    logp = jax.nn.log_softmax(z)
    nll = -jnp.take_along_axis(logp, y[:, None], axis=1)[:, 0]
    return jnp.sum(nll * row_weight) / jnp.sum(row_weight), jax.nn.softmax(z)

zb = jnp.array([[2.0, 1.0, 0.5], [0.3, 0.2, 0.1]])
yb = jnp.array([0, 0])
valid_cols = jnp.array([[True, True, False], [False, False, False]])   # row 2 is padding: fully masked
row_weight = jnp.array([1.0, 0.0])                                      # float weights: 0 for the padding row

out = {}
for name, fill in [("❌ fill -inf", -jnp.inf), ("✅ fill finfo.min", NEG)]:
    (loss, probs), gz = jax.value_and_grad(masked_ce, has_aux=True)(zb, yb, valid_cols, row_weight, fill)
    out[name] = (loss, probs, gz)
    print(f"{name}: loss = {float(loss)}, gradient finite = {bool(jnp.all(jnp.isfinite(gz)))}, "
          f"probabilities of the padding row = {probs[1]}")

loss_ok, probs_ok, g_ok = out["✅ fill finfo.min"]
ref = -np_log_softmax(np.array([2.0, 1.0]))[0]               # only row 1 and its first two classes are valid
assert np.isclose(float(loss_ok), ref, rtol=1e-6) and bool(jnp.all(jnp.isfinite(g_ok)))
assert np.allclose(probs_ok[1], 1 / 3)                        # fully masked row → uniform distribution, harmless
assert np.isnan(float(out["❌ fill -inf"][0])) and bool(jnp.all(jnp.isnan(out["❌ fill -inf"][1][1])))

def locate_nan(fn, *args, user_funcs=()):
    """Run fn with jax_debug_nans on and return (error message, line in user code). finally makes sure the switch is turned off."""
    import traceback
    jax.config.update("jax_debug_nans", True)
    jax.clear_caches()          # measured in 0.12: after vmap has been used in the process, NaNs are missed unless the cache is cleared (see the next cell)
    try:
        jax.block_until_ready(fn(*args))
        return None, None
    except FloatingPointError as e:
        frames = [f for f in traceback.extract_tb(e.__traceback__) if f.name in user_funcs]
        where_ = f"{frames[-1].name}(): {frames[-1].line.strip()}" if frames else "(inside the backward computation generated by JAX)"
        return str(e).splitlines()[0], where_
    finally:
        jax.config.update("jax_debug_nans", False)

def encoder(x):
    h = x * 2.0
    scale = jnp.log(h - 3.0)                 # when h < 3, log(negative) = NaN
    return jnp.sum(scale * h)

def masked_log(x):
    return jnp.where(x > 0, jnp.log(x), 0.0)  # the forward NaN is dropped by where and never reaches the output

x_bad = jnp.array([1.0, 5.0])
tests = [
    ("eager",                       encoder,                       (x_bad,)),
    ("jit: NaN in output",          jax.jit(encoder),              (x_bad,)),
    ("grad(jit): NaN in backward",  jax.grad(jax.jit(masked_log)), (jnp.array(0.0),)),
    ("jit: NaN not in output",      jax.jit(masked_log),           (jnp.array(-1.0),)),
    ("eager: same function",        masked_log,                    (jnp.array(-1.0),)),
    ("jit: inf only",               jax.jit(lambda x: 1.0 / x),    (jnp.array(0.0),)),
]
results = {}
for name, fn, args in tests:
    msg, loc = locate_nan(fn, *args, user_funcs={"encoder", "masked_log"})
    results[name] = msg
    print(f"[{name}]\n    {'no error' if msg is None else 'FloatingPointError: ' + msg}" + (f"\n    located at → {loc}" if loc else ""))

assert results["eager"] and results["jit: NaN in output"] and results["grad(jit): NaN in backward"]
assert results["jit: NaN not in output"] is None and results["eager: same function"] is not None
assert results["jit: inf only"] is None
assert jax.config.jax_debug_nans is False            # the switch has been restored

# inf has to be caught separately with jax_debug_infs
jax.config.update("jax_debug_infs", True)
jax.clear_caches()
try:
    jax.jit(lambda x: 1.0 / x)(jnp.array(0.0))
    caught_inf = None
except FloatingPointError as e:
    caught_inf = str(e).splitlines()[0]
finally:
    jax.config.update("jax_debug_infs", False)
print("jax_debug_infs:", caught_inf)
assert caught_inf is not None

from jax.experimental import checkify

checked = jax.jit(checkify.checkify(encoder, errors=checkify.float_checks))
err, out = checked(x_bad)
print("checkify report:", err.get())
err_ok, _ = checked(jnp.array([3.0, 5.0]))
print("Normal input:", err_ok.get())
err2, _ = checked(x_bad)                      # second call, hits the cache
print("Bad input again:", err2.get())
assert err.get() is not None and err_ok.get() is None and err2.get() is not None
try:
    err.throw()
except Exception as e:
    print(f"err.throw() → {type(e).__name__}")