Chapter 2: Accelerating Functions with JIT Compilation

# Shared imports and helper functions (every later cell depends on these)
import time
import timeit
import logging
import re
import contextlib
import warnings
from functools import partial

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax

print("JAX version:", jax.__version__)
print("Devices:", jax.devices())
print("Device the array is on:", jnp.ones(1).device)   # 0.12 recommends the x.device attribute


def bench(f, *args, warmup=2, runs=20, **kwargs):
    """Timing helper: warm up first (triggering compilation), then run several times and average.

    Uses jax.block_until_ready to wait until the result is really computed -- it handles any PyTree,
    and also non-JAX return values (returned as is). Returns (mean seconds, standard deviation).
    """
    for _ in range(warmup):
        jax.block_until_ready(f(*args, **kwargs))
    ts = []
    for _ in range(runs):
        t0 = time.perf_counter()
        jax.block_until_ready(f(*args, **kwargs))
        ts.append(time.perf_counter() - t0)
    return float(np.mean(ts)), float(np.std(ts))


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 expect_error(fn, exc_type=Exception):
    """Assert that fn() raises exc_type: print the exception type and first line of its message, and return the exception.
    If nothing is raised, raise AssertionError -- used in the notebook to prove that some code really does fail."""
    try:
        out = fn()
    except exc_type as e:
        first = str(e).strip().splitlines()[0]
        print(f"❌ {type(e).__name__}: {first[:160]}")
        return e
    raise AssertionError(f"Expected {exc_type.__name__} to be raised, but there was no error; result: {out}")


class _CompileLog(logging.Handler):
    """Collects the "Compiling jit(<function name>)" logs produced by jax_log_compiles."""

    def __init__(self, echo):
        super().__init__()
        self.names, self.echo = [], echo

    def emit(self, record):
        m = re.match(r"Compiling jit\((.*?)\)", record.getMessage())
        if m:
            self.names.append(m.group(1))
            if self.echo:
                print(f"    [compile] jit({m.group(1)})", flush=True)

    def count(self, name=None):
        return sum(1 for n in self.names if name is None or n == name)


@contextlib.contextmanager
def count_compiles(echo=True):
    """Context manager: count the XLA compilations that actually happen inside the block (optionally filtered by function name).

    How it works: temporarily turn on jax_log_compiles and attach a collector to the "jax" logger.
    Usage: with count_compiles() as log: ...;  assert log.count("function_name") == k
    """
    h = _CompileLog(echo)
    logger = logging.getLogger("jax")
    logger.addHandler(h)
    old = jax.config.jax_log_compiles
    jax.config.update("jax_log_compiles", True)
    try:
        yield h
    finally:
        jax.config.update("jax_log_compiles", old)
        logger.removeHandler(h)

def many_small_ops(x):
    # 50 tiny element-wise operations: each one is dispatched separately
    for _ in range(50):
        x = x * 1.0001 + 0.001
    return x

x_small = jnp.ones(100)
many_small_ops_jit = jax.jit(many_small_ops)

t_eager, _ = bench(many_small_ops, x_small)
t_jit, _ = bench(many_small_ops_jit, x_small)
print(f"eager: {t_eager*1e6:8.1f} µs")
print(f"jit  : {t_jit*1e6:8.1f} µs   about {t_eager/t_jit:.0f}x faster")
# The jit version must be faster (what it mainly saves is 50 dispatches); timing depends on load, so assert_faster re-measures if it fails
assert_faster(lambda: bench(many_small_ops_jit, x_small)[0], lambda: bench(many_small_ops, x_small)[0])

# Look at the HLO after XLA compilation and find the fusion (fused kernel)
fused = jax.jit(lambda a, x, b: jnp.tanh(a * x + b))
hlo = fused.lower(2.0, jnp.ones(1000), 1.0).compile().as_text()
fusion_lines = [line.strip() for line in hlo.splitlines() if "fusion(" in line]
for line in fusion_lines:
    print(line[:160])
assert len(fusion_lines) == 1   # multiply, add, and tanh leave just one fusion instruction
# Only one fusion instruction: multiply, add, and tanh are merged into a single loop kernel (kind=kLoop)

# Course example: a "slightly complex" computation
def complex_computation(x, weight, bias):
    y = jnp.dot(x, weight) + bias
    z = jnp.tanh(y)
    return jnp.mean(z)

# Generate data: use jax.random.key and split off an independent subkey for each array (old vs. new style is compared later in this section)
key = jax.random.key(0)
kx, kw, kb = jax.random.split(key, 3)
x_data = jax.random.normal(kx, (1000, 500))
weight_data = jax.random.normal(kw, (500, 200))
bias_data = jax.random.normal(kb, (200,))

# Usage 1: decorator
@jax.jit
def compiled_computation_decorator(x, weight, bias):
    y = jnp.dot(x, weight) + bias
    z = jnp.tanh(y)
    return jnp.mean(z)

# Usage 2: higher-order function
compiled_computation_functional = jax.jit(complex_computation)

# Timing procedure (uses the more precise time.perf_counter and always calls block_until_ready)

# Original function: warm up once, then time it
_ = complex_computation(x_data, weight_data, bias_data).block_until_ready()
t0 = time.perf_counter()
result_original = complex_computation(x_data, weight_data, bias_data).block_until_ready()
print(f"Original function time:                         {time.perf_counter() - t0:.6f} s")

# Decorator version: the first call includes tracing + compilation
t0 = time.perf_counter()
result_compiled_decorator = compiled_computation_decorator(x_data, weight_data, bias_data).block_until_ready()
print(f"Compiled (decorator), first call (with compile): {time.perf_counter() - t0:.6f} s")

# Second call: hits the cache
t0 = time.perf_counter()
result_compiled_decorator = compiled_computation_decorator(x_data, weight_data, bias_data).block_until_ready()
print(f"Compiled (decorator), second call (cached):      {time.perf_counter() - t0:.6f} s")

print("Results match:", jnp.allclose(result_original, result_compiled_decorator))
assert jnp.allclose(result_original, result_compiled_decorator)

# Functional version: call once to trigger compilation, then time it
_ = compiled_computation_functional(x_data, weight_data, bias_data).block_until_ready()
t0 = time.perf_counter()
_ = compiled_computation_functional(x_data, weight_data, bias_data).block_until_ready()
print(f"Compiled (functional), later call:               {time.perf_counter() - t0:.6f} s")

# ❌ Old style (outdated): PRNGKey returns an untyped raw uint32 key, and the same key is reused for several arrays
try:
    key_old = jax.random.PRNGKey(0)
    a_old = jax.random.normal(key_old, (6,))
    b_old = jax.random.normal(key_old, (3,))
    print("key =", key_old, " dtype =", key_old.dtype)
    print("a[:3] =", a_old[:3])
    print("b     =", b_old)
    print("The two arrays have the same prefix:", bool(jnp.allclose(a_old[:3], b_old)))
except Exception as e:
    print("Error:", type(e).__name__, e)

# ✅ Recommended: jax.random.key returns a typed key; split off an independent subkey for each use
key = jax.random.key(0)
k_a, k_b = jax.random.split(key)
a_new = jax.random.normal(k_a, (6,))
b_new = jax.random.normal(k_b, (3,))
print("key dtype =", key.dtype)
print("a[:3] =", a_new[:3])
print("b     =", b_new)

assert jnp.issubdtype(key.dtype, jax.dtypes.prng_key)     # a new-style key has a dedicated key dtype
assert not jnp.allclose(a_new[:3], b_new)                 # independent subkeys -> the sequences are no longer the same
# Counter-check: reusing the same key always gives the same prefix (exactly what we want to avoid)
assert jnp.allclose(jax.random.normal(key, (6,))[:3], jax.random.normal(key, (3,)))
print("✅ Assertions passed")

def double(x):
    return x * 2

# Same function object -> shared compilation cache
a_jit = jax.jit(double)
b_jit = jax.jit(double)
a_jit(1.0).block_until_ready()
t0 = time.perf_counter(); b_jit(1.0).block_until_ready()
print(f"First call of another jax.jit(double): {(time.perf_counter()-t0)*1e3:.3f} ms (hits the shared cache)")

# A new lambda on every loop iteration -> recompiles every time
for i in range(3):
    t0 = time.perf_counter()
    jax.jit(lambda x: x * 2)(1.0).block_until_ready()
    print(f"New lambda #{i}: {(time.perf_counter()-t0)*1e3:.3f} ms (compiling every time)")

# Passing jit arguments with partial
@partial(jax.jit, static_argnames=["n"])
def power(x, n):
    return x ** n

print("power(2.0, n=3) =", power(2.0, n=3))
assert power(2.0, n=3) == 8.0

def my_simple_func(x, y):
    a = jnp.sin(x)
    b = jnp.cos(y)
    return a + b

x_example = jnp.ones(3)
y_example = jnp.zeros(3)

print(jax.make_jaxpr(my_simple_func)(x_example, y_example))

# Input values don't appear in the Jaxpr; constants hard-coded in the code appear as literals
jp = str(jax.make_jaxpr(lambda x: x * 2.0)(jnp.array([7.5, 8.5])))
assert "2.0" in jp and "7.5" not in jp and "8.5" not in jp

@jax.jit
def slow_function(x):
    # Loop 5 times (the Python loop is unrolled during tracing)
    for _ in range(5):
        if x.shape[0] == x.shape[1]:
            # Square matrix: matrix multiplication
            x = jnp.dot(x, x.T) + 0.5 * x
        else:
            # Non-square matrix: element-wise operations
            x = x * x + 0.5 * x
    return x

key = jax.random.key(0)
data = jax.random.normal(key, (100, 100))

t0 = time.perf_counter(); result1 = slow_function(data).block_until_ready()
print(f"First call (compile) time:          {time.perf_counter() - t0:.4f} s")

t0 = time.perf_counter(); result2 = slow_function(data).block_until_ready()
print(f"Second call (cached) time:          {time.perf_counter() - t0:.4f} s")

data_different_shape = jax.random.normal(key, (150, 150))
t0 = time.perf_counter(); result3 = slow_function(data_different_shape).block_until_ready()
print(f"Call with a new shape:              {time.perf_counter() - t0:.4f} s")

t0 = time.perf_counter(); result4 = slow_function(data).block_until_ready()
print(f"Call with the original shape again: {time.perf_counter() - t0:.4f} s")

jaxpr = jax.make_jaxpr(slow_function.__wrapped__)(data)  # __wrapped__ gets back the original, un-jitted function
print("Number of dot_general occurrences:", str(jaxpr).count("dot_general"))
assert str(jaxpr).count("dot_general") == 5   # range(5) is unrolled into 5 copies

# Raw compile log: create the inputs first -- jnp.ones itself compiles small kernels in eager mode, which would get mixed into the log
a3, b3, a4 = jnp.ones(3), jnp.ones(3) * 7, jnp.ones(4)
g = jax.jit(lambda x: x * 2)
jax.config.update("jax_log_compiles", True)
print("Call 1: f32[3]", flush=True); g(a3)               # compiles
print("Call 2: f32[3], different values", flush=True); g(b3)      # cache hit, no log
print("Call 3: f32[4]", flush=True); g(a4)               # new shape -> compiles
jax.config.update("jax_log_compiles", False)

# E1 shape: changing only values doesn't recompile; changing the shape does
def e1(x):
    return x * 2

x3a, x3b, x4 = jnp.ones(3), jnp.arange(3.0), jnp.ones(4)   # create the inputs first
e1_jit = jax.jit(e1)
with count_compiles() as log:
    e1_jit(x3a)
    e1_jit(x3b)                 # same shape, different values
    n_same_shape = log.count("e1")
    e1_jit(x4)                  # new shape
assert n_same_shape == 1
assert log.count("e1") == 2
print("E1 ✅ Same shape compiles only once; after changing the shape it has compiled", log.count("e1"), "times in total")

# E2 dtype: int32 recompiles; with the default config float64 is truncated to float32 and doesn't recompile
def e2(x):
    return x + 1

xf32 = jnp.ones(3, dtype=jnp.float32)
xi32 = jnp.ones(3, dtype=jnp.int32)
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    x64_req = jnp.ones(3, dtype=jnp.float64)     # x64 is off by default: truncated
assert x64_req.dtype == jnp.float32 and len(w) >= 1
print("Requested float64, actually got:", x64_req.dtype, "| warning:", str(w[0].message)[:60], "...")

e2_jit = jax.jit(e2)
with count_compiles() as log:
    e2_jit(xf32)
    e2_jit(x64_req)
    n_after_f64 = log.count("e2")
    e2_jit(xi32)
assert n_after_f64 == 1        # float32 and "float64" share one compiled result
assert log.count("e2") == 2    # only int32 triggers a second compile

# For comparison: temporarily turn on x64 with a context manager, and float64 becomes a genuinely different dtype (doesn't affect later cells)
def e2_x64(x):
    return x + 1

with jax.enable_x64(True):
    y32 = jnp.ones(3, dtype=jnp.float32)
    y64 = jnp.ones(3, dtype=jnp.float64)
    assert y64.dtype == jnp.float64
    e2_x64_jit = jax.jit(e2_x64)
    with count_compiles() as log64:
        e2_x64_jit(y32)
        e2_x64_jit(y64)
assert log64.count("e2_x64") == 2
print("E2 ✅ Default config: float32/\"float64\" compile 1 time in total, 2 with int32; with x64 on, float32 and float64 each compile once")

# E3 weak types: Python scalars are "weakly typed" and have a different signature from strongly typed arrays of the same dtype
def e3(x):
    return x + 1

e3_jit = jax.jit(e3)
x_strong = jnp.float32(1.0)
with count_compiles() as log:
    e3_jit(1.0)
    e3_jit(2.0)          # Python float: weak f32[], only the value changed -> hit
    e3_jit(x_strong)     # strong f32[] -> new signature
    e3_jit(1)            # Python int: weak i32[] -> new signature
assert log.count("e3") == 3

# Arrays "derived" from Python scalars are often weakly typed too
for name, arr in [("jnp.array(5.0)", jnp.array(5.0)), ("jnp.full(3, 7.0)", jnp.full(3, 7.0)),
                  ("jnp.ones(3)", jnp.ones(3)), ("jnp.float32(1.0)", jnp.float32(1.0))]:
    print(f"{name:18s} weak_type={arr.weak_type}")
assert jnp.array(5.0).weak_type and jnp.full(3, 7.0).weak_type
assert not jnp.ones(3).weak_type and not jnp.float32(1.0).weak_type
print("E3 ✅ Weak/strong types and float/int scalars compiled", log.count("e3"), "times in total")

# E4 static argument values: each new value compiles once; values already seen hit the cache
def e4(x, n):
    for _ in range(n):
        x = jnp.sin(x)
    return x

e4_jit = jax.jit(e4, static_argnames="n")
xe4 = jnp.ones(3)
with count_compiles() as log:
    for n in [2, 3, 2, 3]:
        e4_jit(xe4, n=n)
assert log.count("e4") == 2
print("E4 ✅ n = 2,3,2,3 compiled", log.count("e4"), "times in total")

# E5 PyTree structure: keys, container type, and length are all part of the structure
def e5(d):
    return d["a"] * 2

def e5_seq(t):
    return t[0] * 2

d1 = {"a": jnp.ones(3)}
d2 = {"a": jnp.zeros(3)}                     # same structure, different values
d3 = {"a": jnp.ones(3), "b": jnp.ones(3)}    # one extra key
xs5 = jnp.ones(3)
e5_jit, e5_seq_jit = jax.jit(e5), jax.jit(e5_seq)
with count_compiles() as log:
    e5_jit(d1)
    e5_jit(d2)
    n_same_tree = log.count("e5")
    e5_jit(d3)
    e5_seq_jit((xs5,))       # tuple
    e5_seq_jit([xs5])        # list: different container type -> new structure
assert n_same_tree == 1 and log.count("e5") == 2
assert log.count("e5_seq") == 2
print("E5 ✅ Adding a dict key -> recompiles; tuple and list -> compile once each")

# E6 positional vs. keyword arguments, and omitting default arguments
def e6(x, t):
    return x * t

e6_jit = jax.jit(e6, static_argnums=(1,))
x6 = jnp.ones(3)
with count_compiles() as log:
    e6_jit(x6, 2.0)          # passed positionally
    e6_jit(x6, t=2.0)        # passed by keyword: also static, but a different cache key
assert log.count("e6") == 2


def e6_default(x, tol=1e-6):
    return x + tol

e6d_jit = jax.jit(e6_default)
print(jax.make_jaxpr(e6_default)(x6))     # when omitted: the default is written into the Jaxpr as a constant (the float32 approximation of 1e-6 shows as 9.99…e-07)
with count_compiles() as log:
    e6d_jit(x6)              # omitted: the default is baked in as a constant
    e6d_jit(x6, 1e-3)        # passed positionally: tol becomes a traced argument -> new entry
    e6d_jit(x6, 1e-4)        # passed positionally with a new value -> hit
    e6d_jit(x6, tol=1e-4)    # passed by keyword -> new entry
assert log.count("e6_default") == 3
print("E6 ✅ Positional and keyword each get their own entry; omitted/positional/keyword compile 3 times in total")

# E7 a new jit object every time: the same function object shares the cache; a new lambda each time compiles every time
def e7(x):
    return x * 2

x7 = jnp.ones(3)
with count_compiles() as log:
    for _ in range(3):
        jax.jit(e7)(x7)                  # same function object e7 -> shared cache
    n_same_fn = log.count("e7")
    for _ in range(3):
        jax.jit(lambda x: x * 2)(x7)     # a new lambda object each time
assert n_same_fn == 1
assert log.count("<lambda>") == 3
print("E7 ✅ jax.jit(e7) ×3 compiles only once; a new lambda ×3 compiles 3 times -> jit once outside the loop")

# E8 global variables: changing one doesn't trigger recompilation -> the old value keeps being used
scale_e8 = 1.0

def e8(x):
    return x * scale_e8

e8_jit = jax.jit(e8)
x8 = jnp.ones(2)
with count_compiles() as log:
    r1 = e8_jit(x8)
    scale_e8 = 5.0               # change the global variable
    r2 = e8_jit(x8)
assert log.count("e8") == 1                 # no recompilation...
assert jnp.allclose(r2, 1.0)                # ...so the result is still computed with the old value 1.0
print("E8 ✅ After changing the global, the result is still", r2, "(for the sneakier case where old and new values coexist, see Problem 5 in 2.6)")

def affine_sin(x):
    return jnp.sin(x) * 2

jf = jax.jit(affine_sin)
xs = jnp.ones(3)

print("== trace → Jaxpr ==")
print(jf.trace(xs).jaxpr)

lowered = jf.lower(xs)
print("== lower → StableHLO ==")
print(lowered.as_text())

t0 = time.perf_counter()
compiled = lowered.compile()
print(f"== compile took {(time.perf_counter()-t0)*1e3:.2f} ms ==")
print("Result:", compiled(xs))
assert jnp.allclose(compiled(xs), jnp.sin(xs) * 2)
ca = compiled.cost_analysis()
print("cost_analysis:", {k: ca[k] for k in ("flops", "transcendentals", "bytes accessed")})

# A Compiled object can't handle a different shape
expect_error(lambda: compiled(jnp.ones(4)), TypeError)

def conditional_func_py(x):
    if x > 0:
        return x * 2
    else:
        return x / 2

print(conditional_func_py(5.0))    # 10.0
print(conditional_func_py(-4.0))   # -2.0

# Try jit: x is replaced by a tracer, so `if x > 0` can't get True/False
jitted_conditional_py = jax.jit(conditional_func_py)
expect_error(lambda: jitted_conditional_py(jnp.array(5.0)), jax.errors.ConcretizationTypeError)

def conditional_func_lax(x):
    def true_branch(val):
        return val * 2

    def false_branch(val):
        return val / 2

    return lax.cond(x > 0, true_branch, false_branch, x)

jitted_conditional_lax = jax.jit(conditional_func_lax)
print(jitted_conditional_lax(jnp.array(5.0)))    # 10.0
print(jitted_conditional_lax(jnp.array(-4.0)))   # -2.0
assert jitted_conditional_lax(jnp.array(5.0)) == 10.0 and jitted_conditional_lax(jnp.array(-4.0)) == -2.0

def sum_first_n_py(arr, n):
    total = 0.0
    for i in range(n):
        total = total + arr[i]
    return total

jitted_sum_first_3 = jax.jit(sum_first_n_py, static_argnums=1)
my_array = jnp.arange(10.0)
print(jitted_sum_first_3(my_array, 3))   # 3.0 = 0 + 1 + 2
assert jitted_sum_first_3(my_array, 3) == 3.0

# See what unrolling looks like: with n=3 there are 3 element reads + 3 additions
print(jax.make_jaxpr(sum_first_n_py, static_argnums=1)(my_array, 3))

def sum_first_n_lax(arr, n):
    def body_fun(i, current_total):
        return current_total + arr[i]

    initial_val = 0.0
    final_total = lax.fori_loop(0, n, body_fun, initial_val)
    return final_total

jitted_sum_lax = jax.jit(sum_first_n_lax)
my_array = jnp.arange(10.0)
print(jitted_sum_lax(my_array, jnp.array(3)))   # 3.0  = 0+1+2
print(jitted_sum_lax(my_array, jnp.array(5)))   # 10.0 = 0+1+2+3+4
assert jitted_sum_lax(my_array, jnp.array(3)) == 3.0 and jitted_sum_lax(my_array, jnp.array(5)) == 10.0

try:
    jax.jit(conditional_func_py)(jnp.array(5.0))
except jax.errors.ConcretizationTypeError as e:
    print("Type:", type(e).__name__)
    print("Inheritance chain:", [c.__name__ for c in type(e).__mro__[:3]])

def loop_n(x, n):
    for _ in range(n):
        x = x + 1
    return x

expect_error(lambda: jax.jit(loop_n)(1.0, 3), jax.errors.TracerIntegerConversionError)

# Check the inheritance: only Bool is a subclass of ConcretizationTypeError; the common parent of all four is JAXTypeError
E = jax.errors
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.TracerBoolConversionError,
            E.TracerIntegerConversionError, E.TracerArrayConversionError))
print("Inheritance verified: to catch them all, use jax.errors.JAXTypeError")

body = lambda i, s: s + i
print("Static bound  → ", "scan" if "scan[" in str(jax.make_jaxpr(lambda s: lax.fori_loop(0, 3, body, s))(0.0)) else "while")
print("Traced bound  → ", "while" if "while[" in str(jax.make_jaxpr(lambda n: lax.fori_loop(0, n, body, 0.0))(3)) else "scan")

# Static bound: differentiable. d/dx x^3 at x=2 = 12
g_static = jax.grad(lambda x: lax.fori_loop(0, 3, lambda i, s: s * x, 1.0))(2.0)
print("grad(static fori):", g_static)
assert g_static == 12.0

# while_loop: reverse-mode differentiation raises an error
expect_error(lambda: jax.grad(lambda x: lax.while_loop(lambda s: s < 10, lambda s: s * x, 1.0))(2.0), ValueError)

branches = [lambda x: x, lambda x: x * 10, lambda x: -x]
sw = jax.jit(lambda i, x: lax.switch(i, branches, x))
print("switch(1):", sw(1, 3.0), "  switch(99) is clamped to the last branch:", sw(99, 3.0))

# while_loop: start at 1 and keep doubling until >= 10, counting the steps
final, steps = lax.while_loop(lambda s: s[0] < 10, lambda s: (s[0] * 2, s[1] + 1), (1, 0))
print("while_loop:", final, "steps", steps)

# scan: running sum
carry, ys = lax.scan(lambda c, x: (c + x, c + x), 0.0, jnp.arange(1.0, 5.0))
print("scan: final carry =", carry, " per-step outputs =", ys)

# Branch outputs with mismatched types raise an error
expect_error(lambda: lax.cond(True, lambda x: x, lambda x: jnp.int32(1), 1.0), TypeError)

relu2 = jax.jit(lambda x: jnp.where(x > 0, x * 2, x))
print("jnp.where element-wise:", relu2(jnp.array([-1.0, 2.0])))
assert jnp.array_equal(relu2(jnp.array([-1.0, 2.0])), jnp.array([-1.0, 4.0]))

# Under vmap, a lax.cond with a batched predicate becomes select_n: both branches are computed
jaxpr_v = str(jax.make_jaxpr(jax.vmap(lambda x: lax.cond(x > 0, lambda v: v * 2, lambda v: -v, x)))(jnp.ones(3)))
print("select_n in the Jaxpr of vmap(cond):", "select_n" in jaxpr_v, "| cond still there:", "cond[" in jaxpr_v)
assert "select_n" in jaxpr_v and "cond[" not in jaxpr_v

def unrolled(x, n):
    for _ in range(n):
        x = jnp.sin(x) + 1
    return x

def rolled(x, n):
    return lax.fori_loop(0, n, lambda i, x: jnp.sin(x) + 1, x)

for n in [10, 100, 1000]:
    t0 = time.perf_counter(); jax.jit(unrolled, static_argnums=1).lower(1.0, n).compile(); t_u = time.perf_counter() - t0
    t0 = time.perf_counter(); jax.jit(rolled, static_argnums=1).lower(1.0, n).compile(); t_r = time.perf_counter() - t0
    print(f"n={n:5d}  unrolled: {t_u*1e3:7.1f} ms   fori_loop: {t_r*1e3:6.1f} ms")

def conditional_double(x, threshold):
    if x > threshold:
        return x * 2
    else:
        return x

# Without marking static arguments: both x and threshold are traced, so the if can't get a concrete True/False
jitted_conditional_double = jax.jit(conditional_double)
expect_error(lambda: jitted_conditional_double(jnp.array(5.0), threshold=0.0), jax.errors.ConcretizationTypeError)

# ❌ Marking only threshold as static: instead of printing 10.0 and -2.0, it raises an error
print("static_argnums=(1,)      →", end=" ")
expect_error(lambda: jax.jit(conditional_double, static_argnums=(1,))(jnp.array(5.0), threshold=0.0),
             jax.errors.ConcretizationTypeError)
print("static_argnames=threshold →", end=" ")
expect_error(lambda: jax.jit(conditional_double, static_argnames=("threshold",))(jnp.array(-2.0), threshold=0.0),
             jax.errors.ConcretizationTypeError)

# ✅ Fix 1: the condition depends on the value of x -> lax.cond or jnp.where (threshold can be traced, and changing it doesn't recompile)
@jax.jit
def conditional_double_cond(x, threshold):
    return lax.cond(x > threshold, lambda v: v * 2, lambda v: v, x)

@jax.jit
def conditional_double_where(x, threshold):
    return jnp.where(x > threshold, x * 2, x)

for fn in (conditional_double_cond, conditional_double_where):
    assert fn(jnp.array(5.0), 0.0) == 10.0
    assert fn(jnp.array(-2.0), 0.0) == -2.0
print("Fix 1 ✅ lax.cond:", conditional_double_cond(jnp.array(5.0), 0.0), conditional_double_cond(jnp.array(-2.0), 0.0))

# ✅ Fix 2: static_argnames only really works when the if condition depends entirely on static arguments
@partial(jax.jit, static_argnames=["double"])
def maybe_double(x, double):
    if double:            # double is a Python bool, known at trace time
        return x * 2
    return x

assert maybe_double(jnp.array(5.0), double=True) == 10.0
assert maybe_double(jnp.array(5.0), double=False) == 5.0
print("Fix 2 ✅ static switch:", maybe_double(jnp.array(5.0), double=True), maybe_double(jnp.array(5.0), double=False))

def conditional_double_static(x, threshold):
    return jnp.where(x > threshold, x * 2, x)

jitted_func = jax.jit(conditional_double_static, static_argnames=("threshold",))
x5, x1 = jnp.array(5.0), jnp.array(1.0)      # create the inputs first
with count_compiles() as log:
    print("Call 1 (x=5, threshold=0.0)", flush=True);  r_a = jitted_func(x5, threshold=0.0)    # compiles
    print("Call 2 (x=1, threshold=0.0)", flush=True);  r_b = jitted_func(x1, threshold=0.0)    # reuses
    print("Call 3 (x=5, threshold=10.0)", flush=True); r_c = jitted_func(x5, threshold=10.0)   # recompiles
    print("Call 4 (x=1, threshold=10.0)", flush=True); r_d = jitted_func(x1, threshold=10.0)   # reuses
assert log.count("conditional_double_static") == 2
assert [float(r) for r in (r_a, r_b, r_c, r_d)] == [10.0, 2.0, 5.0, 1.0]
print("✅ 4 calls compiled", log.count("conditional_double_static"), "times in total")

from dataclasses import dataclass

# 1) Arrays and lists can't be static arguments; tuples can
expect_error(lambda: jax.jit(lambda x, t: x * t, static_argnums=1)(1.0, jnp.array(2.0)), ValueError)
expect_error(lambda: jax.jit(lambda x, shape: x.reshape(shape), static_argnums=1)(jnp.arange(6.0), [2, 3]), ValueError)
reshaped = jax.jit(lambda x, shape: x.reshape(shape), static_argnums=1)(jnp.arange(6.0), (2, 3))
assert reshaped.shape == (2, 3)

# 2) Custom classes as static arguments: hashing by identity vs. by value
class CfgById:                     # default __hash__/__eq__: by object identity
    def __init__(self, k):
        self.k = k

@dataclass(frozen=True)
class CfgByValue:                  # frozen dataclass: hashes by field values
    k: int

def scale_by_id(x, cfg):
    return x * cfg.k

def scale_by_value(x, cfg):
    return x * cfg.k

xc = jnp.ones(3)
j_id = jax.jit(scale_by_id, static_argnums=1)
j_val = jax.jit(scale_by_value, static_argnums=1)
with count_compiles(echo=False) as log:
    j_id(xc, CfgById(2)); j_id(xc, CfgById(2))          # same contents, but two objects
    j_val(xc, CfgByValue(2)); j_val(xc, CfgByValue(2))  # equal values -> cache hit
assert log.count("scale_by_id") == 2
assert log.count("scale_by_value") == 1
print("Config class hashed by identity compiled", log.count("scale_by_id"), "times; frozen dataclass compiled", log.count("scale_by_value"), "time(s)")

# 4) Shape information is static
@jax.jit
def sum_rows_or_all(x):
    return x.sum() if x.ndim == 1 else x.sum(axis=0)

assert sum_rows_or_all(jnp.ones(3)) == 3.0
assert jnp.array_equal(sum_rows_or_all(jnp.ones((2, 3))), jnp.full(3, 2.0))
print("Python if based on ndim ✅", sum_rows_or_all(jnp.ones(3)), sum_rows_or_all(jnp.ones((2, 3))))

@jax.jit
def condition_on_value(x):
    if x > 0:
        return x * 2
    else:
        return x / 2

expect_error(lambda: condition_on_value(jnp.array(5.0)), jax.errors.ConcretizationTypeError)

# Fix: lax.cond(predicate, true branch, false branch, operand)
@jax.jit
def condition_on_value_lax(x):
    def true_fun(operand):
        return operand * 2

    def false_fun(operand):
        return operand / 2

    return lax.cond(x > 0, true_fun, false_fun, x)

print(condition_on_value_lax(jnp.array(5.0)))    # 10.0
print(condition_on_value_lax(jnp.array(-4.0)))   # -2.0
assert condition_on_value_lax(jnp.array(5.0)) == 10.0 and condition_on_value_lax(jnp.array(-4.0)) == -2.0

def variable_loop(x, n):
    total = x
    for i in range(n):
        total = total + i
    return total

# n is a traced value (a JAX array), so range(n) can't determine the number of iterations
expect_error(lambda: jax.jit(variable_loop)(10.0, jnp.array(3)), jax.errors.TracerIntegerConversionError)

# Fix 1: lax.fori_loop (here combined with static_argnums to make n static)
def fixed_loop_fori(x, n_static):
    def body_fun(i, current_total):
        return current_total + i
    return lax.fori_loop(0, n_static, body_fun, x)

jitted_fori = jax.jit(fixed_loop_fori, static_argnums=(1,))
print(jitted_fori(jnp.array(10.0), 3))   # 13.0 = 10 + 0 + 1 + 2
assert jitted_fori(jnp.array(10.0), 3) == 13.0

# Fix 2: lax.scan (a loop with state; here a running sum)
@jax.jit
def cumulative_sum_scan(xs):
    def scan_op(carry, x):
        new_carry = carry + x
        return new_carry, new_carry   # (carry for the next step, output of this step)
    _, ys = lax.scan(scan_op, 0.0, xs)
    return ys

print(cumulative_sum_scan(jnp.array([1.0, 2.0, 3.0, 4.0])))   # [ 1.  3.  6. 10.]
assert jnp.array_equal(cumulative_sum_scan(jnp.array([1.0, 2.0, 3.0, 4.0])), jnp.array([1.0, 3.0, 6.0, 10.0]))

# ✅ Pass n as a traced value: the result is correct, and different n share one compiled result
def loop_add(x, n):
    return lax.fori_loop(0, n, lambda i, t: t + i, x)

loop_add_jit = jax.jit(loop_add)
x10 = jnp.array(10.0)
with count_compiles() as log:
    results = [float(loop_add_jit(x10, n)) for n in (3, 5, 100)]
print("n=3, 5, 100 →", results)
assert results == [13.0, 20.0, 10.0 + sum(range(100))]
assert log.count("loop_add") == 1
print("✅ Three different values of n compiled only", log.count("loop_add"), "time")

@jax.jit
def function_with_side_effect(x):
    print(f"TRACE: Running Python code for x={x}")   # runs only during tracing
    time.sleep(0.1)                                    # also runs only during tracing
    return x * x

print(function_with_side_effect(jnp.array(3.0)))   # prints TRACE (x is a tracer)
print(function_with_side_effect(jnp.array(4.0)))   # no TRACE, and no sleep
print(function_with_side_effect(jnp.array(5.0)))

@jax.jit
def debug_demo(x):
    print("Plain print (tracing only):", x)
    jax.debug.print("jax.debug.print (every run): x = {x}", x=x)
    return x * x

debug_demo(jnp.array(3.0)).block_until_ready()
debug_demo(jnp.array(4.0)).block_until_ready()

# disable_jit: x is a concrete value inside the function, so a Python if works too
with jax.disable_jit():
    print("Under disable_jit:", condition_on_value(jnp.array(5.0)))

@jax.jit
def process_data(x):
    return jnp.sum(x * 2.0)

inputs = [jnp.ones(s) for s in [(3,), (3,), (4,), (4,)]]   # create the inputs first so only process_data itself is timed
with count_compiles(echo=False) as log:
    for x in inputs:
        t0 = time.perf_counter()
        process_data(x).block_until_ready()
        print(f"shape={x.shape}: {time.perf_counter() - t0:.4f} s")
assert log.count("process_data") == 2     # once for (3,), once for (4,)
print("✅ Compiled", log.count("process_data"), "times in total")

@jax.jit
def masked_mean(x_padded, length):
    # length is a traced value: use a mask instead of the slice x[:length] (slice lengths must be static)
    mask = jnp.arange(x_padded.shape[0]) < length
    return jnp.sum(jnp.where(mask, x_padded, 0.0)) / length

def bucket_size(n):
    return max(8, 1 << (n - 1).bit_length())   # round up to a power of 2, minimum 8

def mean_any_length(x):
    n = x.shape[0]
    padded = jnp.pad(x, (0, bucket_size(n) - n))
    return masked_mean(padded, n)

lengths = [3, 5, 7, 9, 12, 16, 20]
data_by_len = [jnp.arange(1.0, n + 1) for n in lengths]    # create the inputs first
with count_compiles(echo=False) as log:
    for x in data_by_len:
        assert jnp.allclose(mean_any_length(x), x.mean())
assert log.count("masked_mean") == 3
print("All 7 lengths give correct results; masked_mean compiled only", log.count("masked_mean"), "times (buckets: 8, 16, 32)")

learning_rate = 0.01

@jax.jit
def update_weights(params, grads):
    return params - learning_rate * grads   # learning_rate is frozen in during tracing

params = jnp.array([1.0, 2.0])
grads = jnp.array([0.5, -0.1])

r1 = update_weights(params, grads)
print(r1)                              # [0.995 2.001]
learning_rate = 1000.0                 # change the global variable
r2 = update_weights(params, grads)
print(r2)                              # still [0.995 2.001]: it uses the old value 0.01
assert np.allclose(r1, [0.995, 2.001]) and np.allclose(r2, r1)