Chapter 7: Custom Kernels with Pallas

# Shared setup: imports, version check, cross-checking and timing helpers
import sys, time, re, logging, functools, inspect, contextlib
import numpy as np
import jax
import jax.numpy as jnp
from jax import lax

print("Python", sys.version.split()[0], "| JAX", jax.__version__, "| backend", jax.default_backend(), jax.devices())
assert jax.__version__.startswith("0.12")
assert jax.default_backend() == "cpu", "This chapter runs on CPU in interpret mode (on Colab, use a CPU runtime)"

def f64(a):
    """Convert to NumPy float64, for high-precision reference implementations"""
    return np.asarray(a, dtype=np.float64)

def max_err(a, ref):
    """Max absolute error (both converted to float64 before comparing)"""
    return float(np.max(np.abs(f64(a) - f64(ref))))

def bench(f, *args, n=5):
    """Average time (ms): warm up once, then loop n times. Timings are only printed, never asserted"""
    jax.block_until_ready(f(*args))
    t0 = time.perf_counter()
    for _ in range(n):
        out = f(*args)
    jax.block_until_ready(out)
    return (time.perf_counter() - t0) / n * 1e3

@contextlib.contextmanager
def quiet_callback_logs():
    """When the TPU simulator reports an error through io_callback, it first logs a full traceback; silence it while demonstrating ❌ code"""
    lg = logging.getLogger("jax._src.callback")
    old = lg.disabled
    lg.disabled = True
    try:
        yield
    finally:
        lg.disabled = old

rng = np.random.default_rng(0)   # only for a little NumPy-side data; JAX randomness always uses jax.random.key

# Probe 1: import paths and key names
from jax.experimental import pallas as pl
from jax.experimental.pallas import tpu as pltpu            # TPU (Mosaic) specific: VMEM, CompilerParams, InterpretParams...
from jax.experimental.pallas import triton as pltriton      # GPU · Triton backend
from jax.experimental.pallas import mosaic_gpu as plgpu     # GPU · Mosaic GPU backend

must_have = ["pallas_call", "BlockSpec", "GridSpec", "program_id", "num_programs", "when",
             "cdiv", "ds", "CompilerParams", "CostEstimate", "debug_print", "Squeezed"]
for name in must_have:
    assert hasattr(pl, name), name
assert hasattr(pltpu, "InterpretParams") and hasattr(pltpu, "CompilerParams") and hasattr(pltpu, "VMEM")
assert hasattr(plgpu, "InterpretGPUParams") and hasattr(pltriton, "CompilerParams")

# Names from older tutorials: gone in 0.12
gone = {"pl.load": hasattr(pl, "load"), "pl.store": hasattr(pl, "store"),
        "pl.InterpretParams": hasattr(pl, "InterpretParams"),
        "pltpu.TPUCompilerParams": hasattr(pltpu, "TPUCompilerParams")}
print("Names that no longer exist:", gone)
assert not any(gone.values())

sig = inspect.signature(pl.pallas_call)
print("pallas_call parameters:", list(sig.parameters))
print("interpret default:", sig.parameters["interpret"].default)
assert sig.parameters["interpret"].default is False
assert {"out_shape", "grid", "in_specs", "out_specs", "scratch_shapes", "compiler_params",
        "input_output_aliases", "cost_estimate", "interpret", "name"} <= set(sig.parameters)

# Probe 2: what does the interpret argument actually accept? All three forms tested on CPU
def add_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] + y_ref[...]

x = jnp.arange(8, dtype=jnp.float32)
y = jnp.full((8,), 10.0, jnp.float32)
sds = jax.ShapeDtypeStruct(x.shape, x.dtype)

results = {}
for label, interp in [("interpret=True", True),
                      ("interpret=pltpu.InterpretParams()", pltpu.InterpretParams()),
                      ("interpret=plgpu.InterpretGPUParams()", plgpu.InterpretGPUParams())]:
    out = pl.pallas_call(add_kernel, out_shape=sds, interpret=interp)(x, y)
    results[label] = np.asarray(out)
    print(f"{label:40s} -> {results[label]}")
    np.testing.assert_array_equal(results[label], f64(x) + f64(y))

# ❌ Wrong: no interpret on the CPU backend — the CPU has no Pallas compilation backend
try:
    pl.pallas_call(add_kernel, out_shape=sds)(x, y)
    err = None
except Exception as e:
    err = e
    print(type(e).__name__, ":", e)
assert isinstance(err, ValueError) and "Only interpret mode is supported on CPU" in str(err)

# ✅ Correct: use one global switch on this machine; on GPU/TPU just change it to False
INTERPRET = True
out = pl.pallas_call(add_kernel, out_shape=sds, interpret=INTERPRET)(x, y)
np.testing.assert_array_equal(out, f64(x) + f64(y))
print(out)

# ❌ Wrong: reading/writing Refs with pl.load / pl.store as in older tutorials — both were removed from the pl namespace in 0.12
def add_kernel_old(x_ref, y_ref, o_ref):
    pl.store(o_ref, ..., pl.load(x_ref, ...) + pl.load(y_ref, ...))
try:
    pl.pallas_call(add_kernel_old, out_shape=sds, interpret=INTERPRET)(x, y)
    err = None
except AttributeError as e:
    err = e
    print(type(e).__name__, ":", e)
assert isinstance(err, AttributeError)

# ✅ Correct: read and write Refs with NumPy-style indexing; for a dynamic slice use pl.ds(start, size)
def add_kernel_halves(x_ref, y_ref, o_ref):
    o_ref[pl.ds(0, 4)] = x_ref[pl.ds(0, 4)] + y_ref[pl.ds(0, 4)]   # first half
    o_ref[4:] = x_ref[4:] + y_ref[4:]                              # second half: static slices work too
out = pl.pallas_call(add_kernel_halves, out_shape=sds, interpret=INTERPRET)(x, y)
np.testing.assert_array_equal(out, f64(x) + f64(y))
print(out)

# 1) The pallas_call primitive in the jaxpr
add_call = pl.pallas_call(add_kernel, out_shape=sds, interpret=INTERPRET)
cj = jax.make_jaxpr(add_call)(x, y)
prims = [str(e.primitive) for e in cj.jaxpr.eqns]
print("Primitives in the outer jaxpr:", prims)
inner = cj.jaxpr.eqns[0].params["jaxpr"]          # pallas_call has the kernel's own jaxpr nested inside
print("The kernel's jaxpr:\n", inner)
assert prims == ["pallas_call"]
assert all("Ref" in str(v.aval) for v in inner.invars)   # all of the kernel's arguments are Refs

# 2) Interpret mode: the compiled HLO has only ordinary ops and a while loop, no custom-call
hlo_interp = jax.jit(add_call).lower(x, y).compile().as_text()
print("Number of custom-calls in the interpret-mode HLO:", hlo_interp.count("custom-call"))
assert "custom-call" not in hlo_interp

# 3) Cross-lowering to TPU (not executed! we only look at the generated StableHLO)
add_tpu = pl.pallas_call(add_kernel, out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32))  # interpret off
tpu_text = jax.jit(add_tpu).trace(jax.ShapeDtypeStruct((8, 128), jnp.float32),
                                  jax.ShapeDtypeStruct((8, 128), jnp.float32)).lower(lowering_platforms=("tpu",)).as_text()
line = [l.strip() for l in tpu_text.splitlines() if "tpu_custom_call" in l][0]
print("TPU lowering result:", line[:110], "...")
assert "tpu_custom_call" in tpu_text

# Multiple outputs: out_shape is a tuple, and the kernel takes one more output Ref
def sum_diff_kernel(x_ref, y_ref, s_ref, d_ref):
    x, y = x_ref[...], y_ref[...]          # read them first to avoid reading twice
    s_ref[...] = x + y
    d_ref[...] = x - y

kx, ky = jax.random.split(jax.random.key(0))
a = jax.random.normal(kx, (4, 6))
b = jax.random.normal(ky, (4, 6))
s, d = pl.pallas_call(sum_diff_kernel,
                      out_shape=(jax.ShapeDtypeStruct(a.shape, a.dtype),
                                 jax.ShapeDtypeStruct(a.shape, a.dtype)),
                      interpret=INTERPRET)(a, b)
np.testing.assert_allclose(f64(s), f64(a) + f64(b), rtol=1e-6)
np.testing.assert_allclose(f64(d), f64(a) - f64(b), rtol=1e-6)
print("Multi-output kernel is correct:", s.shape, d.shape)

# ❌ Wrong: the kernel returns its result (like an ordinary function) instead of writing it into the output Ref
def double_kernel_bad(x_ref, o_ref):
    return x_ref[...] * 2
try:
    pl.pallas_call(double_kernel_bad, out_shape=jax.ShapeDtypeStruct(a.shape, a.dtype), interpret=INTERPRET)(a)
    err = None
except ValueError as e:
    err = e
    print(type(e).__name__, ": … the kernel function", re.search(r"should return None.*", str(e)).group(0))
assert err is not None and "should return None" in str(err)

# ✅ Correct: write into o_ref; the function itself returns nothing
def double_kernel(x_ref, o_ref):
    o_ref[...] = x_ref[...] * 2
out = pl.pallas_call(double_kernel, out_shape=jax.ShapeDtypeStruct(a.shape, a.dtype), interpret=INTERPRET)(a)
np.testing.assert_allclose(f64(out), 2 * f64(a), rtol=1e-7)
print("✅", out.shape)

# ❌ Wrong: the kernel captures a JAX array (say, a weight) through a closure
W_bias = jnp.arange(6, dtype=jnp.float32)
def bias_kernel_bad(x_ref, o_ref):
    o_ref[...] = x_ref[...] + W_bias          # W_bias is an external array constant
try:
    pl.pallas_call(bias_kernel_bad, out_shape=jax.ShapeDtypeStruct(a.shape, a.dtype), interpret=INTERPRET)(a)
    err = None
except ValueError as e:
    err = e
    print(type(e).__name__, ": … the kernel function", re.search(r"captures constants.*", str(e)).group(0))
assert err is not None and "captures constants" in str(err)

# ✅ Correct: always pass arrays in as inputs (Python scalars may be captured by closure; they're treated as compile-time constants)
scale = 0.5                                   # Python float: allowed
def bias_kernel(x_ref, w_ref, o_ref):
    o_ref[...] = x_ref[...] * scale + w_ref[...]
out = pl.pallas_call(bias_kernel, out_shape=jax.ShapeDtypeStruct(a.shape, a.dtype), interpret=INTERPRET)(a, W_bias)
np.testing.assert_allclose(f64(out), f64(a) * 0.5 + f64(W_bias), rtol=1e-6)
print("✅ Weights passed in as an argument; result is correct")

def add2d(x, y, bm=64, bn=128):
    """2-D elementwise addition: each program handles one (bm, bn) block"""
    M, N = x.shape
    spec = pl.BlockSpec((bm, bn), lambda i, j: (i, j))     # program (i, j) takes block (i, j)
    return pl.pallas_call(add_kernel,                        # same kernel as in §0: it only cares about "one block"
                          out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
                          grid=(M // bm, N // bn),
                          in_specs=[spec, spec], out_specs=spec,
                          interpret=INTERPRET)(x, y)

kx, ky = jax.random.split(jax.random.key(1))
X = jax.random.normal(kx, (256, 512))
Y = jax.random.normal(ky, (256, 512))
Z = add2d(X, Y)
np.testing.assert_allclose(f64(Z), f64(X) + f64(Y), rtol=1e-6)
print("Tiled addition with grid=(4, 4) is correct, output shape", Z.shape)

# Have each program write its own (program_id(0), program_id(1)) into the block it owns, to see the grid → block mapping
def pid_kernel(o_ref):
    i, j = pl.program_id(0), pl.program_id(1)
    o_ref[...] = jnp.full(o_ref.shape, 10 * i + j, o_ref.dtype)   # tens digit is i, ones digit is j

pid_map = pl.pallas_call(pid_kernel, out_shape=jax.ShapeDtypeStruct((4, 6), jnp.int32),
                         grid=(2, 3), out_specs=pl.BlockSpec((2, 2), lambda i, j: (i, j)),
                         interpret=INTERPRET)()
print(pid_map)
expected = np.kron(np.array([[0, 1, 2], [10, 11, 12]]), np.ones((2, 2), int))
np.testing.assert_array_equal(pid_map, expected)

# ❌ Wrong: index_map returns an "element offset" instead of a "block index" — no error, silently wrong results
xs = jnp.arange(64, dtype=jnp.float32)
bad = pl.pallas_call(double_kernel, out_shape=jax.ShapeDtypeStruct(xs.shape, xs.dtype), grid=(4,),
                     in_specs=[pl.BlockSpec((16,), lambda i: (i * 16,))],    # wrong: i*16 is an element offset
                     out_specs=pl.BlockSpec((16,), lambda i: (i,)),
                     interpret=INTERPRET)(xs)
wrong = ~np.isclose(f64(bad), 2 * f64(xs))
print("Wrong elements:", int(wrong.sum()), "/ 64; outputs 16–19:", np.asarray(bad[16:20]))
assert wrong.sum() > 0      # in interpret mode an out-of-range block index is clamped to the last block, so the result is quietly wrong

# ✅ Correct: index_map returns block indices, and Pallas multiplies by block_shape to get the offset
good = pl.pallas_call(double_kernel, out_shape=jax.ShapeDtypeStruct(xs.shape, xs.dtype), grid=(4,),
                      in_specs=[pl.BlockSpec((16,), lambda i: (i,))],
                      out_specs=pl.BlockSpec((16,), lambda i: (i,)),
                      interpret=INTERPRET)(xs)
np.testing.assert_array_equal(good, 2 * f64(xs))
print("✅", np.asarray(good[16:20]))

def matmul_rowcol(a, b, bm=64, bn=128):
    """A matmul that doesn't tile along K: A takes (bm, K) row blocks, B takes (K, bn) column blocks"""
    M, K = a.shape
    _, N = b.shape
    def kernel(a_ref, b_ref, o_ref):
        o_ref[...] = jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((M, N), jnp.float32),
                          grid=(M // bm, N // bn),
                          in_specs=[pl.BlockSpec((bm, K), lambda i, j: (i, 0)),   # row block i; K has only one block
                                    pl.BlockSpec((K, bn), lambda i, j: (0, j))],  # column block j
                          out_specs=pl.BlockSpec((bm, bn), lambda i, j: (i, j)),
                          interpret=INTERPRET)(a, b)

ka, kb = jax.random.split(jax.random.key(2))
A = jax.random.normal(ka, (256, 384))
B = jax.random.normal(kb, (384, 512))
C_rc = matmul_rowcol(A, B)
ref_AB = f64(A) @ f64(B)                       # NumPy float64 reference
print("Max absolute error:", max_err(C_rc, ref_AB))
np.testing.assert_allclose(f64(C_rc), ref_AB, rtol=1e-4, atol=1e-4)

visit_order = []
def recorder(token, grid_point, core_id):
    visit_order.append(tuple(int(c) for c in grid_point))
    return token                               # must return the token unchanged, or you get "Mismatched number of outputs"

def copy_kernel(x_ref, o_ref):
    o_ref[...] = x_ref[...]

spec = pl.BlockSpec((8, 128), lambda i, j: (i, j))
xg = jax.random.normal(jax.random.key(3), (16, 384))
yg = pl.pallas_call(copy_kernel, out_shape=jax.ShapeDtypeStruct(xg.shape, xg.dtype), grid=(2, 3),
                    in_specs=[spec], out_specs=spec,
                    interpret=pltpu.InterpretParams(grid_point_recorder=recorder))(xg)
jax.block_until_ready(yg)
print("Visit order:", visit_order)
np.testing.assert_array_equal(yg, xg)
assert visit_order == [(i, j) for i in range(2) for j in range(3)]     # lexicographic order

x100 = jnp.arange(1, 101, dtype=jnp.float32)          # length 100, block size 32 → 4 blocks; the last block has only 4 valid elements
peek = pl.pallas_call(copy_kernel, out_shape=jax.ShapeDtypeStruct((4, 32), jnp.float32), grid=(pl.cdiv(100, 32),),
                      in_specs=[pl.BlockSpec((32,), lambda i: (i,))],
                      out_specs=pl.BlockSpec((None, 32), lambda i: (i, 0)),   # None: this dim is squeezed out; the kernel sees (32,)
                      interpret=INTERPRET)(x100)
tail = np.asarray(peek[3])
print("What the 4th block read:", tail[:8], "...")
assert np.array_equal(tail[:4], [97, 98, 99, 100]) and np.isnan(tail[4:]).all()

# But an elementwise kernel is unaffected: the out-of-bounds part is discarded on write
dbl = pl.pallas_call(double_kernel, out_shape=jax.ShapeDtypeStruct(x100.shape, x100.dtype), grid=(pl.cdiv(100, 32),),
                     in_specs=[pl.BlockSpec((32,), lambda i: (i,))], out_specs=pl.BlockSpec((32,), lambda i: (i,)),
                     interpret=INTERPRET)(x100)
np.testing.assert_array_equal(dbl, 2 * f64(x100))
print("The elementwise kernel is still correct on a non-divisible length")

kx, ky = jax.random.split(jax.random.key(4))
Xb = jax.random.normal(kx, (1024, 1024))
Yb = jax.random.normal(ky, (1024, 1024))
add_jnp = jax.jit(lambda x, y: x + y)
add_pl_1 = jax.jit(functools.partial(add2d, bm=1024, bn=1024))   # 1 block
add_pl_64 = jax.jit(functools.partial(add2d, bm=128, bn=128))    # 64 blocks
np.testing.assert_allclose(f64(add_pl_64(Xb, Yb)), f64(Xb) + f64(Yb), rtol=1e-6)
print(f"jnp + XLA                   : {bench(add_jnp, Xb, Yb):8.3f} ms")
print(f"Pallas interpret, 1 block   : {bench(add_pl_1, Xb, Yb):8.3f} ms")
print(f"Pallas interpret, 64 blocks : {bench(add_pl_64, Xb, Yb):8.3f} ms   ← more blocks, longer interpretation loop (printed only, no conclusion)")

def matmul_acc(a, b, bm=64, bk=128, bn=128):
    """Tiled matmul accumulating along K, float32 output, using o_ref directly as the accumulator"""
    M, K = a.shape
    _, N = b.shape
    def kernel(a_ref, b_ref, o_ref):
        @pl.when(pl.program_id(2) == 0)            # first K step of each output block: zero it
        def _init():
            o_ref[...] = jnp.zeros_like(o_ref)
        o_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((M, N), jnp.float32),
                          grid=(M // bm, N // bn, K // bk),       # reduction dim k goes last
                          in_specs=[pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)),
                                    pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))],
                          out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),
                          interpret=INTERPRET)(a, b)

C_acc = matmul_acc(A, B)                    # A: (256, 384), B: (384, 512) → grid (4, 4, 3)
print("Max absolute error (vs. float64):", max_err(C_acc, ref_AB))
np.testing.assert_allclose(f64(C_acc), ref_AB, rtol=1e-4, atol=1e-4)

# ❌ Wrong: testing program_id with a Python if — it's a tracer at tracing time and can't be converted to bool
def mm_kernel_pyif(a_ref, b_ref, o_ref):
    if pl.program_id(2) == 0:
        o_ref[...] = jnp.zeros_like(o_ref)
    o_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
try:
    pl.pallas_call(mm_kernel_pyif, out_shape=jax.ShapeDtypeStruct((256, 512), jnp.float32), grid=(4, 4, 3),
                   in_specs=[pl.BlockSpec((64, 128), lambda i, j, k: (i, k)),
                             pl.BlockSpec((128, 128), lambda i, j, k: (k, j))],
                   out_specs=pl.BlockSpec((64, 128), lambda i, j, k: (i, j)), interpret=INTERPRET)(A, B)
    err = None
except jax.errors.TracerBoolConversionError as e:
    err = e
    print(type(e).__name__, ":", str(e).splitlines()[0])
assert isinstance(err, jax.errors.TracerBoolConversionError)

# ✅ Correct: decorate a no-argument function with @pl.when(condition) (exactly what matmul_acc does)
def mm_kernel_when(a_ref, b_ref, o_ref):
    @pl.when(pl.program_id(2) == 0)
    def _init():
        o_ref[...] = jnp.zeros_like(o_ref)
    o_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
C_when = pl.pallas_call(mm_kernel_when, out_shape=jax.ShapeDtypeStruct((256, 512), jnp.float32), grid=(4, 4, 3),
                        in_specs=[pl.BlockSpec((64, 128), lambda i, j, k: (i, k)),
                                  pl.BlockSpec((128, 128), lambda i, j, k: (k, j))],
                        out_specs=pl.BlockSpec((64, 128), lambda i, j, k: (i, j)), interpret=INTERPRET)(A, B)
np.testing.assert_allclose(f64(C_when), ref_AB, rtol=1e-4, atol=1e-4)
print("✅ The pl.when version is correct, max error", max_err(C_when, ref_AB))

# ❌ Wrong: forgetting to zero at k == 0 and just doing +=
def mm_kernel_noinit(a_ref, b_ref, o_ref):
    o_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
mm_noinit = lambda interp: pl.pallas_call(
    mm_kernel_noinit, out_shape=jax.ShapeDtypeStruct((256, 512), jnp.float32), grid=(4, 4, 3),
    in_specs=[pl.BlockSpec((64, 128), lambda i, j, k: (i, k)), pl.BlockSpec((128, 128), lambda i, j, k: (k, j))],
    out_specs=pl.BlockSpec((64, 128), lambda i, j, k: (i, j)), interpret=interp)(A, B)
try:
    r_nan = mm_noinit(INTERPRET)                                          # uninitialized memory = NaN
    r_zero = mm_noinit(pltpu.InterpretParams(uninitialized_memory="zero")) # even if the initial memory happens to be 0...
    print("interpret=True: fraction of NaN =", float(np.isnan(f64(r_nan)).mean()))
    print("TPU simulator (uninitialized memory filled with 0): max error =", max_err(r_zero, ref_AB))
except Exception as e:
    print(type(e).__name__, e)
    raise
assert np.isnan(f64(r_nan)).all()
assert max_err(r_zero, ref_AB) > 1.0     # the output buffer is reused across blocks, so the previous block's result gets accumulated in

# ✅ Correct: zero explicitly at k == 0 (matmul_acc). Correct under both interpreters and both fill values for uninitialized memory
for label, interp in [("interpret=True", INTERPRET), ("InterpretParams()", pltpu.InterpretParams()),
                      ('InterpretParams(uninitialized_memory="zero")', pltpu.InterpretParams(uninitialized_memory="zero"))]:
    C_ok = pl.pallas_call(mm_kernel_when, out_shape=jax.ShapeDtypeStruct((256, 512), jnp.float32), grid=(4, 4, 3),
                          in_specs=[pl.BlockSpec((64, 128), lambda i, j, k: (i, k)),
                                    pl.BlockSpec((128, 128), lambda i, j, k: (k, j))],
                          out_specs=pl.BlockSpec((64, 128), lambda i, j, k: (i, j)), interpret=interp)(A, B)
    np.testing.assert_allclose(f64(C_ok), ref_AB, rtol=1e-4, atol=1e-4)
    print(f"✅ {label:46s} max error {max_err(C_ok, ref_AB):.2e}")

def matmul_bf16(a, b, bm=64, bk=128, bn=128):
    """bf16 inputs/outputs, accumulation in a float32 scratch"""
    M, K = a.shape
    _, N = b.shape
    def kernel(a_ref, b_ref, o_ref, acc_ref):         # the last argument is the scratch
        k = pl.program_id(2)
        @pl.when(k == 0)
        def _init():
            acc_ref[...] = jnp.zeros_like(acc_ref)
        acc_ref[...] += jnp.dot(a_ref[...], b_ref[...], preferred_element_type=jnp.float32)
        @pl.when(k == pl.num_programs(2) - 1)        # last K step: write out
        def _store():
            o_ref[...] = acc_ref[...].astype(o_ref.dtype)
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((M, N), a.dtype),
                          grid=(M // bm, N // bn, K // bk),
                          in_specs=[pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)),
                                    pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))],
                          out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),
                          scratch_shapes=[pltpu.VMEM((bm, bn), jnp.float32)],
                          interpret=INTERPRET)(a, b)

A16, B16 = A.astype(jnp.bfloat16), B.astype(jnp.bfloat16)
C16 = matmul_bf16(A16, B16)
ref16 = f64(A16) @ f64(B16)                # reference: the same bf16 inputs, computed in float64
rel = np.max(np.abs(f64(C16) - ref16) / (np.abs(ref16) + 1.0))
print("Output dtype:", C16.dtype, "| max relative error:", rel, "(bf16's unit roundoff is about 2^-8 ≈ 3.9e-3)")
assert C16.dtype == jnp.bfloat16
assert rel < 2 ** -7

def compute_dtype(dt):
    """Compute precision inside the block: promote bf16/f16 to float32; keep float32 and float64 as they are (a blunt astype(float32) would lower float64's precision)"""
    return jnp.promote_types(dt, jnp.float32)

def softmax_kernel(x_ref, o_ref):
    x = x_ref[...].astype(compute_dtype(x_ref.dtype))
    m = jnp.max(x, axis=-1, keepdims=True)          # numerically stable: subtract the row max first
    e = jnp.exp(x - m)
    o_ref[...] = (e / jnp.sum(e, axis=-1, keepdims=True)).astype(o_ref.dtype)

def row_call(kernel, x, *extra, br=8):
    """Call a kernel over row blocks: each block is br rows × full width; extra are column-broadcast arguments (shape (1, C), shared by all blocks)"""
    R, Cc = x.shape
    row_spec = pl.BlockSpec((br, Cc), lambda i: (i, 0))
    vec_spec = pl.BlockSpec((1, Cc), lambda i: (0, 0))      # every program takes the same block (block 0)
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
                          grid=(pl.cdiv(R, br),),
                          in_specs=[row_spec] + [vec_spec] * len(extra), out_specs=row_spec,
                          interpret=INTERPRET)(x, *extra)

def softmax_ref64(x):
    x = f64(x); e = np.exp(x - x.max(-1, keepdims=True)); return e / e.sum(-1, keepdims=True)

xs_sm = jax.random.normal(jax.random.key(5), (50, 256)) * 3       # 50 rows: not divisible by 8
sm = row_call(softmax_kernel, xs_sm)
print("softmax max error:", max_err(sm, softmax_ref64(xs_sm)))
np.testing.assert_allclose(f64(sm), softmax_ref64(xs_sm), rtol=1e-5, atol=1e-7)
np.testing.assert_allclose(f64(sm).sum(-1), 1.0, rtol=1e-5)

# ❌ Wrong: softmax without subtracting the max — for large inputs exp overflows to inf, and inf/inf = NaN
def softmax_kernel_naive(x_ref, o_ref):
    e = jnp.exp(x_ref[...])
    o_ref[...] = e / jnp.sum(e, axis=-1, keepdims=True)
big = xs_sm + 100.0                                     # softmax is shift-invariant, so the correct answer doesn't change
try:
    sm_bad = row_call(softmax_kernel_naive, big)
    print("Fraction of NaN:", float(np.isnan(f64(sm_bad)).mean()))
except Exception as e:
    print(type(e).__name__, e); raise
assert np.isnan(f64(sm_bad)).any()

# ✅ Correct: after subtracting the row max, the result shifted by 100 matches the original
sm_big = row_call(softmax_kernel, big)
np.testing.assert_allclose(f64(sm_big), softmax_ref64(xs_sm), rtol=1e-5, atol=1e-7)
print("✅ Result unchanged after a shift of 100, max error", max_err(sm_big, softmax_ref64(xs_sm)))

EPS = 1e-5
def layernorm_kernel(x_ref, g_ref, b_ref, o_ref):
    x = x_ref[...].astype(compute_dtype(x_ref.dtype))
    mu = jnp.mean(x, axis=-1, keepdims=True)
    xc = x - mu
    var = jnp.mean(xc * xc, axis=-1, keepdims=True)        # two-pass: center first, then take the mean of the squares
    o_ref[...] = (xc * lax.rsqrt(var + EPS) * g_ref[...] + b_ref[...]).astype(o_ref.dtype)

def layernorm_ref64(x, g, b):
    x = f64(x); mu = x.mean(-1, keepdims=True); var = ((x - mu) ** 2).mean(-1, keepdims=True)
    return (x - mu) / np.sqrt(var + EPS) * f64(g) + f64(b)

kx, kg, kb = jax.random.split(jax.random.key(6), 3)
xln = jax.random.normal(kx, (50, 256))
gamma = 1.0 + 0.1 * jax.random.normal(kg, (1, 256))        # weights passed in as arguments
beta = 0.1 * jax.random.normal(kb, (1, 256))
ln = row_call(layernorm_kernel, xln, gamma, beta)
print("LayerNorm max error:", max_err(ln, layernorm_ref64(xln, gamma, beta)))
np.testing.assert_allclose(f64(ln), layernorm_ref64(xln, gamma, beta), rtol=1e-4, atol=1e-5)

# ❌ Wrong: variance as E[x²] − E[x]². When the data has a large offset (1e4 here), float32 suffers catastrophic cancellation
def layernorm_kernel_onepass(x_ref, g_ref, b_ref, o_ref):
    x = x_ref[...]
    mu = jnp.mean(x, axis=-1, keepdims=True)
    var = jnp.mean(x * x, axis=-1, keepdims=True) - mu * mu          # subtracting two numbers of order 1e8
    o_ref[...] = (x - mu) * lax.rsqrt(var + EPS) * g_ref[...] + b_ref[...]
x_off = xln + 1e4
ref_off = layernorm_ref64(f64(xln) + 1e4, gamma, beta)          # float64 reference
try:
    ln_bad = row_call(layernorm_kernel_onepass, x_off, gamma, beta)
    err_bad = np.nanmax(np.abs(f64(ln_bad) - ref_off)) if not np.isnan(f64(ln_bad)).all() else np.inf
    print("One-pass max error:", err_bad, "| NaN count:", int(np.isnan(f64(ln_bad)).sum()))
except Exception as e:
    print(type(e).__name__, e); raise
assert err_bad > 0.1 or np.isnan(f64(ln_bad)).any()

# ✅ Correct: two-pass (center first)
ln_good = row_call(layernorm_kernel, x_off, gamma, beta)
err_good = max_err(ln_good, ref_off)
print("Two-pass max error:", err_good)
assert err_good < 5e-3     # the remaining error comes from representing 1e4 + x itself in float32 (an ulp of about 1e-3)

def matmul_masked(a, b, bm=64, bk=64, bn=64, mask_a=True, mask_b=True):
    """Tiled matmul for any shape: the grid rounds up with cdiv, and the K direction is masked"""
    M, K = a.shape
    _, N = b.shape
    def kernel(a_ref, b_ref, o_ref):
        k = pl.program_id(2)
        @pl.when(k == 0)
        def _init():
            o_ref[...] = jnp.zeros_like(o_ref)
        a_blk, b_blk = a_ref[...], b_ref[...]
        if mask_a:   # this if tests a Python bool (a build-time constant), not a tracer, so it's fine
            ka = k * bk + lax.broadcasted_iota(jnp.int32, a_blk.shape, 1)   # columns of the A block = global K index
            a_blk = jnp.where(ka < K, a_blk, 0.0)
        if mask_b:
            kb_ = k * bk + lax.broadcasted_iota(jnp.int32, b_blk.shape, 0)  # rows of the B block = global K index
            b_blk = jnp.where(kb_ < K, b_blk, 0.0)
        o_ref[...] += jnp.dot(a_blk, b_blk, preferred_element_type=jnp.float32)
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((M, N), jnp.float32),
                          grid=(pl.cdiv(M, bm), pl.cdiv(N, bn), pl.cdiv(K, bk)),
                          in_specs=[pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)),
                                    pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))],
                          out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),
                          interpret=INTERPRET)(a, b)

ka, kb = jax.random.split(jax.random.key(7))
Ao = jax.random.normal(ka, (100, 200))        # M=100, K=200, N=130: none divisible by 64
Bo = jax.random.normal(kb, (200, 130))
ref_o = f64(Ao) @ f64(Bo)

# ❌ Wrong: no mask along K, or masking only A — NaN pollutes every output
for label, ma, mb in [("no mask", False, False), ("mask A only", True, False)]:
    try:
        r = matmul_masked(Ao, Bo, mask_a=ma, mask_b=mb)
        print(f"{label:8s}: fraction of NaN = {float(np.isnan(f64(r)).mean()):.2f}")
    except Exception as e:
        print(type(e).__name__, e); raise
    assert np.isnan(f64(r)).all()

# ✅ Correct: mask both A's columns and B's rows with jnp.where
C_o = matmul_masked(Ao, Bo)
print("Shape", C_o.shape, "| max error:", max_err(C_o, ref_o))
np.testing.assert_allclose(f64(C_o), ref_o, rtol=1e-4, atol=1e-4)

def rowsum(x, br=32, bc=128, reduce_last=True, interpret=INTERPRET):
    """Row sums: x (R, C) -> (R,); reduce_last=False demonstrates the wrong grid order"""
    R, Cc = x.shape
    red_axis = 1 if reduce_last else 0
    def kernel(x_ref, o_ref):
        j = pl.program_id(red_axis)
        @pl.when(j == 0)
        def _init():
            o_ref[...] = jnp.zeros_like(o_ref)
        col = j * bc + lax.broadcasted_iota(jnp.int32, x_ref.shape, 1)
        o_ref[...] += jnp.sum(jnp.where(col < Cc, x_ref[...], 0.0), axis=1)   # mask, then sum
    if reduce_last:
        grid, xmap, omap = (pl.cdiv(R, br), pl.cdiv(Cc, bc)), (lambda i, j: (i, j)), (lambda i, j: (i,))
    else:   # wrong order: reduction dim first
        grid, xmap, omap = (pl.cdiv(Cc, bc), pl.cdiv(R, br)), (lambda j, i: (i, j)), (lambda j, i: (i,))
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((R,), jnp.float32), grid=grid,
                          in_specs=[pl.BlockSpec((br, bc), xmap)], out_specs=pl.BlockSpec((br,), omap),
                          compiler_params=pltpu.CompilerParams(dimension_semantics=("parallel", "arbitrary")) if reduce_last else None,
                          interpret=interpret)(x)

xr = jax.random.normal(jax.random.key(8), (96, 300))     # 300 columns: not divisible by 128
rs = rowsum(xr)
print("Row-sum max error:", max_err(rs, f64(xr).sum(1)))
np.testing.assert_allclose(f64(rs), f64(xr).sum(1), rtol=1e-5, atol=1e-4)

# ❌ Wrong: the reduction dim is the first grid dim — the same output block is visited "non-consecutively"
r_true = rowsum(xr, reduce_last=False, interpret=True)     # generic interpreter: happens to be right, hiding the problem
print("Error under interpret=True:", max_err(r_true, f64(xr).sum(1)), "← right by luck")
assert max_err(r_true, f64(xr).sum(1)) < 1e-3
try:
    with quiet_callback_logs():
        jax.block_until_ready(rowsum(xr, reduce_last=False, interpret=pltpu.InterpretParams()))
    err = None
except Exception as e:
    err = e
    msg = [l for l in str(e).splitlines() if "Revisited" in l][0]
    print("TPU simulator:", type(e).__name__, "→", msg.strip())
finally:
    pltpu.reset_tpu_interpret_mode_state()                  # must reset after the simulator raises
assert err is not None and "Revisited block" in str(err)

# ✅ Correct: reduction dim last, and the TPU simulator passes too
rs_tpu = rowsum(xr, reduce_last=True, interpret=pltpu.InterpretParams())
np.testing.assert_allclose(f64(rs_tpu), f64(xr).sum(1), rtol=1e-5, atol=1e-4)
print("✅ Error under the TPU simulator:", max_err(rs_tpu, f64(xr).sum(1)))

def logsumexp_rows(x, br=16, bc=128):
    """Row-wise online logsumexp: x (R, C) -> (R,), tiled along C and updated online across blocks"""
    R, Cc = x.shape
    def kernel(x_ref, o_ref, m_ref, l_ref):
        j = pl.program_id(1)
        @pl.when(j == 0)
        def _init():
            m_ref[...] = jnp.full_like(m_ref, -jnp.inf)
            l_ref[...] = jnp.zeros_like(l_ref)
        col = j * bc + lax.broadcasted_iota(jnp.int32, x_ref.shape, 1)
        xb = jnp.where(col < Cc, x_ref[...], -jnp.inf)         # the identity element for max is -inf
        m_old = m_ref[...]
        m_new = jnp.maximum(m_old, jnp.max(xb, axis=1))
        l_ref[...] = l_ref[...] * jnp.exp(m_old - m_new) + jnp.sum(jnp.exp(xb - m_new[:, None]), axis=1)
        m_ref[...] = m_new
        @pl.when(j == pl.num_programs(1) - 1)
        def _store():
            o_ref[...] = m_ref[...] + jnp.log(l_ref[...])
    return pl.pallas_call(kernel, out_shape=jax.ShapeDtypeStruct((R,), jnp.float32),
                          grid=(pl.cdiv(R, br), pl.cdiv(Cc, bc)),
                          in_specs=[pl.BlockSpec((br, bc), lambda i, j: (i, j))],
                          out_specs=pl.BlockSpec((br,), lambda i, j: (i,)),
                          scratch_shapes=[pltpu.VMEM((br,), jnp.float32), pltpu.VMEM((br,), jnp.float32)],
                          interpret=INTERPRET)(x)

from scipy.special import logsumexp as sp_logsumexp
xl = jax.random.normal(jax.random.key(9), (40, 1000)) * 4 + 50.0   # long rows with large values; 1000 columns not divisible by 128
lse = logsumexp_rows(xl)
ref_lse = sp_logsumexp(f64(xl), axis=1)
print("Online logsumexp max error:", max_err(lse, ref_lse))
np.testing.assert_allclose(f64(lse), ref_lse, rtol=1e-5)
np.testing.assert_allclose(f64(lse), f64(jax.nn.logsumexp(xl, axis=1)), rtol=1e-5)