Tapes and events

Tapes and events#

Record what a stepping kernel saw at every step, and turn its stop rule into a smooth function a runtime can differentiate.

Time: ~8 min · Runs on: CPU · You need: hawk and NumPy; read first: Your own derivative, Stop when you’re done

A gradient through a whole rollout needs, at every step, the state that step read. Keeping every state of every step is the simplest answer and the most expensive one; a launching runtime instead keeps a few checkpoints and, one window of steps at a time, re-runs the window while recording a tape. hawk builds the kernel that does that re-run: the taping twin.

Colab or Kaggle: run this first cell (it installs the packages and downloads the helper files this notebook uses). Elsewhere it only checks the imports.

Hide code cell source

 
# Colab setup, generated by docs/tools/nb_colab_setup.py (do not edit).
import os

try:
    import google.colab  # noqa: F401
    HOSTED = True  # a Colab runtime
except ImportError:
    HOSTED = "KAGGLE_KERNEL_RUN_TYPE" in os.environ  # a Kaggle kernel

if HOSTED:
    import importlib.metadata, pathlib, re, shutil, subprocess, sys, urllib.request

    gpu = bool(shutil.which("nvidia-smi")) and subprocess.run(["nvidia-smi", "-L"], capture_output=True).returncode == 0
    extra = ""
    if gpu:
        # the CUDA major the runtime already has: a CuPy that is installed, else torch's, else the driver's
        major = None
        for m in ("13", "12"):
            try:
                importlib.metadata.version(f"cupy-cuda{m}x")
                major = m
                break
            except importlib.metadata.PackageNotFoundError:
                pass
        if major is None:
            try:
                import torch
                major = "13" if str(torch.version.cuda).startswith("13") else None
            except ImportError:
                pass
        if major is None:
            smi = re.search(r"CUDA Version: (\d+)", subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout)
            major = "13" if smi and int(smi.group(1)) >= 13 else "12"
        extra = f"[cuda{major}]"
        print(f"GPU found: installing the {extra} extras to match this runtime.")
    pre = ["--pre"] if os.environ.get("RAPTOR_PRE") else []  # optional: test a release candidate
    ref = os.environ.get("RAPTOR_REF", "main")  # optional: fetch the helper files from this git ref
    if pre or "RAPTOR_REF" in os.environ:
        print(f"Pre-release test: installing with --pre, helper files from {ref}.")
    # our packages with --upgrade, so a runtime that kept an older release gets the current one
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade", *pre, f"raptor-hawk{extra}"])
    if not gpu:
        print("Running on the CPU (no GPU in this runtime). For a GPU: " + ("Settings → Accelerator → GPU T4 x2." if "KAGGLE_KERNEL_RUN_TYPE" in os.environ else "Runtime → Change runtime type → T4 GPU."))

try:
    import hawk  # noqa: F401
except ImportError as err:
    raise ImportError(
        "hawk not found. Install with `pip install raptor-hawk`. For a GPU, add [cuda12] (or [cuda13]) to raptor-hawk."
    ) from err

import pathlib, tempfile
import numpy as np
import hawk
import hawk.math
from hawk import Mutable, Param, Scalar, Terminated


@hawk.kernel
def ball(dt: Param, terminated: Terminated,
         y: Mutable[Scalar], v: Mutable[Scalar], t: Mutable[Scalar]):
    v_start = v
    v = v_start - 9.81 * dt
    y = y + dt * v_start
    tn = t + dt
    t = tn
    terminated = y < 0.0      # it hit the ground


ball

<Kernel ball slots=7>

The twin#

hawk.diff.tape(kernel, window) returns <name>_tape: the same step, plus three planes it reserves.

plane

kind

what it holds

applied_steps

integer state

the steps this sample has taken (the finishing step counts)

tape

output, window × w per sample

row applied_steps mod window = the state before that step

tape_on

integer parameter

record only when non-zero

The row width w is not chosen: it is every state plane the step both reads and overwrites — exactly what its one-step vjp needs as input.

from hawk.diff import tape

window = 8
ball_tape = tape(ball, window)
work_dir = pathlib.Path(tempfile.mkdtemp())
bundle = hawk.build([ball, ball_tape], work_dir, targets=("host",))
sidecar = next(a.sidecar for a in bundle.artifacts if a.name == "ball_tape")
print(ball_tape.name)
print(sidecar["arg_spec"])

ball_tape
[['mutable', 'applied_steps'], ['mutable', 't'], ['mutable', 'v'], ['mutable', 'y'], ['out', 'tape'], ['lookup', 'finished_count'], ['lookup', 'fused_steps'], ['terminated', 'terminated'], ['uniform', 'dt'], ['uniform', 'tape_on']]

The taped planes are t, v, y (the order of the kernel’s own signature listing), so a row is 3 numbers. Let us run four balls from different heights, one step per call, with the tape on:

def drop(name, tape_on):
    n = 4
    y = np.array([0.3, 0.5, 0.8, 1.0])
    v, t = np.zeros(n), np.zeros(n)
    stopped = np.zeros(n, dtype=bool)
    planes = dict(y=y, v=v, t=t, terminated=stopped)
    if name == "ball_tape":
        planes.update(applied_steps=np.zeros(n, dtype=np.int64),
                      tape=np.full((window * 3, n), np.nan), tape_on=tape_on)
    kernel = hawk.load(work_dir, name)
    history = []
    while not stopped.all():
        history.append(np.stack([t, v, y]).copy())
        hawk.run(kernel, dt=0.05, **planes)
    return planes, history


taped, history = drop("ball_tape", tape_on=1)
print("steps taken:", taped["applied_steps"])

steps taken: [ 6  7  9 10]

Row j of a sample is its state before step j. For the first ball (fewer steps than the window, so no row is reused):

rows = taped["tape"][:, 0].reshape(window, 3)
for j in range(int(taped["applied_steps"][0])):
    print(j, rows[j], "  before step", j, ":", history[j][:, 0])

0 [0.  0.  0.3]   before step 0 : [0.  0.  0.3]
1 [ 0.05   -0.4905  0.3   ]   before step 1 : [ 0.05   -0.4905  0.3   ]
2 [ 0.1      -0.981     0.275475]   before step 2 : [ 0.1      -0.981     0.275475]
3 [ 0.15     -1.4715    0.226425]   before step 3 : [ 0.15     -1.4715    0.226425]
4 [ 0.2     -1.962    0.15285]   before step 4 : [ 0.2     -1.962    0.15285]
5 [ 0.25    -2.4525   0.05475]   before step 5 : [ 0.25    -2.4525   0.05475]

Off is the same code#

With tape_on = 0 the twin runs the same compiled entry and stores nothing: a forward that only needs the stop step (applied_steps) and the final state takes exactly the result the recording re-run will reproduce, bit for bit, because it is the same machine code.

quiet, _ = drop("ball_tape", tape_on=0)
same = all(np.array_equal(quiet[k], taped[k]) for k in ("y", "v", "t", "applied_steps"))
print("state identical:", same, "| tape untouched:", bool(np.isnan(quiet["tape"]).all()))

state identical: True | tape untouched: True

The stop rule as a smooth guard#

terminated = y < 0.0 is a flag; its derivative is zero wherever it exists. To ask when the ball lands, a runtime needs the smooth function that changes sign at the landing: the guard g. hawk.diff.event.event_guard reads it out of the comparison, with one convention — the sample stops where g <= 0:

stop rule

guard

a < b, a <= b

g = a - b

a > b, a >= b

g = b - a

lor(c1, c2)

one guard per comparison

The guard is an ordinary expression of the step’s inputs, so guard_kernel packages it as a kernel jvp and vjp differentiate like any other:

from hawk.diff import jvp
from hawk.diff.event import event_guard, guard_kernel

print(event_guard(ball).rules)
ball_guard = guard_kernel(ball)
dg = hawk.Kernel("ball_guard_jvp", jvp(ball_guard, wrt=("y", "v")))
gdir = pathlib.Path(tempfile.mkdtemp())
hawk.build([ball_guard, dg], gdir, targets=("host",))

y, v, t = np.array([0.3, 0.02]), np.array([0.0, -1.0]), np.zeros(2)
g = np.zeros(2)
hawk.run(hawk.load(gdir, "ball_guard"), dt=0.05, y=y, v=v, t=t,
         terminated=np.zeros(2, dtype=bool), event_g0=g)
dot_g = np.zeros(2)
hawk.run(hawk.load(gdir, "ball_guard_jvp"), dt=0.05, y=y, v=v, t=t,
         terminated=np.zeros(2, dtype=bool), dot_y=np.zeros(2), dot_v=np.ones(2),
         dot_event_g0=dot_g)
print("g        :", g, " (the second ball stops this step: g <= 0)")
print("dg / dv  :", dot_g, " (= dt: the step moves y by dt * v)")

('((y + (dt * v)) < 0.0)',)
g        : [ 0.3  -0.03]  (the second ball stops this step: g <= 0)
dg / dv  : [0.05 0.05]  (= dt: the step moves y by dt * v)

Rules that are refused#

A guard has to be a surface the state crosses. A flag, a count, an equality, an and of two conditions or a rounding on the way to the comparison has no smooth crossing, and is refused by name instead of quietly giving a zero:

@hawk.kernel
def bounded(dt: Param, terminated: Terminated, y: Mutable[Scalar]):
    y = y - dt
    terminated = hawk.math.land(y < 0.0, y > -1.0)


try:
    event_guard(bounded)
except hawk.HawkError as err:
    print(err)

events="refine": the stop rule of 'bounded' is not a smooth crossing ((((y - dt) < 0.0) and ((y - dt) > -1.0))); refine needs g(x) that changes sign, not a count or a flag

Time is part of the state#

Where a stop happens depends on when as much as on where: a moving target, a fixed horizon. Naming the time plane and the step on the kernel makes time a state component the twin can reason about:

@hawk.kernel(time="t", time_step="dt")

The plain kernel does not change. hawk checks once, at decoration, that the body advances t as t = t + dt (a body without a t gets one minted by the twin). The twin then

  • reads dt as a per-sample plane, so every sample may step differently, and

  • writes t in index form, t = t_0 + k * dt with k the steps applied and t_0 a per-sample start-time plane — one multiply and one add (a single fused multiply-add where the compiler contracts it), however long the run, instead of a sum that drifts by an ulp here and there at every step. Below, the twin’s t sits within one rounding of t_0 + k * dt, while adding dt step by step lands somewhere else.

@hawk.kernel(time="t", time_step="dt")
def target(dt: Param, a: Param, terminated: Terminated,
           y: Mutable[Scalar], v: Mutable[Scalar], t: Mutable[Scalar]):
    v_start = v
    v = v_start - 9.81 * dt
    y = y + dt * v_start
    tn = t + dt
    t = tn
    terminated = y <= a * tn      # a surface rising at speed a


target_tape = tape(target, 64)
tdir = pathlib.Path(tempfile.mkdtemp())
hawk.build([target_tape], tdir, targets=("host",))

n = 3
dt = np.array([0.01, 0.02, 0.03])
t0 = np.array([0.0, 0.5, 1.0])
p = dict(y=np.full(n, 1.0), v=np.zeros(n), t=t0.copy(),
         terminated=np.zeros(n, dtype=bool), applied_steps=np.zeros(n, dtype=np.int64),
         tape=np.zeros((64 * 3, n)), tape_on=0, dt=dt, t_0=t0, a=0.2)
while not p["terminated"].all():
    hawk.run(hawk.load(tdir, "target_tape"), **p)
print("steps :", p["applied_steps"])
index = t0 + p["applied_steps"] * dt
summed = t0.copy()
for i in range(n):
    for _ in range(int(p["applied_steps"][i])):
        summed[i] += dt[i]
print("t - index form, in ulps :", (p["t"] - index) / np.spacing(index))
print("summed t - t, in ulps   :", (summed - p["t"]) / np.spacing(index))

steps : [44 21 14]
t - index form, in ulps : [0. 1. 0.]
summed t - t, in ulps   : [4. 3. 2.]

Refining the stop#

A step that crosses the surface overshoots it: the sample stopped somewhere inside its last step. To find where, a runtime re-runs that step with a partial size tau and looks for the tau at which the guard is zero. hawk.diff.refine_entry(kernel) builds <name>_refine for exactly this: one step of size <dt>_tau from the state before the stop step, with time written as (t_0 + k * dt) + tau — so the derivative of time with respect to tau is exactly one. It reads the state as plain inputs and writes <plane>_refined and the guard event_g0, so jvp with respect to tau gives the slope a Newton iteration needs:

from hawk.diff import refine_entry

refine = refine_entry(target)
slope = hawk.Kernel("target_refine_jvp", jvp(refine, wrt=("dt_tau",)))
rdir = pathlib.Path(tempfile.mkdtemp())
hawk.build([refine, slope], rdir, targets=("host",))

state = dict(y=np.array([0.05]), v=np.array([-2.0]), dt=np.array([0.05]),
             t_0=np.zeros(1), applied_steps=np.array([10]), a=0.2)
outs = {k: np.zeros(1) for k in ("dot_y_refined", "dot_v_refined",
                                 "dot_t_refined", "dot_event_g0")}
hawk.run(hawk.load(rdir, "target_refine_jvp"), dt_tau=np.array([0.02]),
         dot_dt_tau=np.ones(1), **state, **outs)
print("d t   / d tau :", outs["dot_t_refined"])
print("d g   / d tau :", outs["dot_event_g0"], " (= v - a: the closing speed)")

d t   / d tau : [1.]
d g   / d tau : [-2.2]  (= v - a: the closing speed)

Stopping on the grid, not one step late#

A stop rule can land exactly on a step: a fixed horizon t > T with T a whole number of steps away. Rounding then decides whether the comparison fires on that step or the next, and a strict > fires a whole step late whenever t equals T exactly. Under refinement the twin replaces the kernel’s own comparison by a closed, tolerant one, g <= 4 * spacing(|t|) * |dg/dt| — a few ulps of the time, scaled by how fast the guard moves with time (spacing is one ulp, hawk.math.spacing). Ask for it with events="refine":

@hawk.kernel(time="t", time_step="dt")
def horizon(dt: Param, T: Scalar, terminated: Terminated, t: Mutable[Scalar]):
    tn = t + dt
    t = tn
    terminated = tn > T


def stop_step(events):
    hdir = pathlib.Path(tempfile.mkdtemp())
    hawk.build([tape(horizon, 4, events=events)], hdir, targets=("host",))
    h, m = np.array([0.1, 0.25, 0.3]), np.array([10, 4, 7])
    p = dict(t=np.zeros(3), terminated=np.zeros(3, dtype=bool),
             applied_steps=np.zeros(3, dtype=np.int64), tape=np.zeros((4, 3)),
             tape_on=0, dt=h, t_0=np.zeros(3), T=m * h)
    while not p["terminated"].all():
        hawk.run(hawk.load(hdir, "horizon_tape"), **p)
    return p["applied_steps"]


print("grid step m        :", [10, 4, 7])
print("kernel's own test  :", stop_step("step"))
print("events='refine'    :", stop_step("refine"))

grid step m        : [10, 4, 7]
kernel's own test  : [11  5  8]
events='refine'    : [10  4  7]

Changing a parameter without rebuilding#

A Param is handed to the kernel by value, with each launch. A runtime that records a launch and replays it would have to record it again to change one. hawk.diff.param_planes(kernel) returns <name>_params, the same kernel with each Param read from a plane of ONE number instead. The plane keeps the parameter’s own name, so you rebind k=... as before; changing it later is a copy of one number, and nothing is recompiled or re-recorded.

Taking the derivative of that kernel gives the gradient with respect to each parameter as a per-sample plane (bar_k, one value per sample), not a single running total. You add the samples up yourself, in the order you choose, so the result is the same every time.

from hawk.diff import param_planes, vjp


@hawk.kernel
def law(x: Scalar, k: Param, y: Mutable[Scalar]):
    y = k * x * x


law_params = param_planes(law)
law_grad = hawk.Kernel("law_grad", vjp(law_params, wrt=("x", "k")))
bundle2 = hawk.build([law_params, law_grad], work_dir, targets=("host",))
print(law_params.name, law_params.walk.arg_spec)
print(next(a.sidecar["param_planes"] for a in bundle2.artifacts
           if a.name == "law_params"))

law_params (('mutable', 'y'), ('per_sample', 'x'), ('lookup', 'k'))
['k']

x = np.array([1.0, 2.0, 3.0])
k = np.array([2.0])          # one number, held in a plane of length 1
y = np.zeros(3)
step = hawk.load(work_dir, "law_params")
hawk.run(step, x=x, k=k, y=y)
print("k = 2:", y)
k[0] = 5.0                   # same built kernel, new value
hawk.run(step, x=x, k=k, y=y)
print("k = 5:", y)

bar_k, bar_x = np.zeros(3), np.zeros(3)
hawk.run(hawk.load(work_dir, "law_grad"), x=x, k=k, bar_y=np.ones(3),
         bar_k=bar_k, bar_x=bar_x)
print("per-sample dJ/dk:", bar_k, " sum:", bar_k.sum())

k = 2: [ 2.  8. 18.]
k = 5: [ 5. 20. 45.]
per-sample dJ/dk: [1. 4. 9.]  sum: 14.0

param_planes(kernel, names) promotes only the parameters you name; the sidecar lists the promoted planes under param_planes. Two limits: a parameter used inside a for loop of the step body is refused, and an integer parameter is promoted but has no gradient (a count has no slope). The plain kernel is never changed.

What just happened#

  • hawk.diff.tape(kernel, window) builds <name>_tape: the step plus applied_steps, a tape of pre-step states and a tape_on switch.

  • The tape width is derived from what the one-step vjp reads; nothing else is stored.

  • tape_on = 0 runs the same compiled entry, so recording and not recording agree bit for bit.

  • event_guard turns a stop comparison into a guard with g <= 0 at the stop, differentiable like any expression; non-smooth rules are refused.

  • hawk.diff.param_planes(kernel) reads each Param from a one-number plane, so a value can change without a rebuild; its vjp gives a per-sample gradient plane.

  • time= / time_step= make time a state component: the twin steps by a per-sample dt and writes t = t_0 + k * dt; refine_entry re-runs the stop step with a partial tau, with d t / d tau = 1.

  • tape(..., events="refine") swaps the stop comparison for a closed test with a few ulps of time tolerance, so a stop on the step grid is taken on its own step.

Try this#

Tape ball with window = 2 and drop it from 1.0: the run is longer than the window, so rows are reused — row j mod 2 holds the most recent step j that landed there.

Next#

eagle’s tutorials drive the twin from the runtime side: recording a rollout, reversing it window by window, and refining the stop time from the guard. Short a primitive instead? See when hawk can’t spell it.