Gradients through conditionals, long runs and events

Gradients through conditionals, long runs and events#

Three things a reverse pass meets once a model stops being a straight line: a step that runs only sometimes, a run too long to keep in memory, and a stop that happens between two steps.

Time: ~12 min · Runs on: CPU (the same cells run on the GPU when the planes are cupy arrays) · You need: eagle, hawk and NumPy; read first: Gradients through a graph: eagle.diff, Gradients through a graph: eagle.diff

Every gradient below is checked against something you can verify by hand: central differences, or a closed form.

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.

Show code cell source

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-eagle{extra}", 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 eagle, hawk  # noqa: F401
except ImportError as err:
    raise ImportError(
        "eagle, hawk not found. Install with `pip install raptor-eagle raptor-hawk`. For a GPU, add [cuda12] (or [cuda13]) to raptor-eagle and raptor-hawk."
    ) from err

A step that runs only sometimes#

eagle.skippable(step, guard) runs a step only while a guard cell is non-zero. The guard is read at every launch, so you flip it between launches and nothing is rebuilt. The reverse pass has to differentiate the launch that happened, so a record=True graph saves the guard cell it saw, and the reverse replays that value. The guard itself is never evaluated again.

Here first computes y = a x², maybe overwrites y with sin x only while the flag is on, and last computes sin y.

import tempfile

import numpy as np
import eagle

import hawk
from hawk import Mutable, Param, Scalar
from hawk import math as m


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


@hawk.kernel
def wave(x: Scalar, y: Mutable[Scalar]):
    y = m.sin(x)


square_plan, wave_plan = eagle.deploy([square, wave], targets=("host",),
                                      cache_dir=tempfile.mkdtemp())
n = 6
x = np.linspace(-1.0, 1.0, n)
flag = np.array([1], dtype="uint32")        # the guard cell: 1 = run `maybe`


def build(x_, a_):
    y, z = np.zeros(n), np.zeros(n)
    maybe = eagle.skippable(wave_plan.bind(x=x_, y=y).launch,
                            eagle.SkipGuard.nonzero(flag))
    g = (eagle.Graph()
         .add(square_plan.bind(x=x_, a=a_, y=y), name="first")
         .add_node(maybe, name="maybe")
         .add(wave_plan.bind(x=y, y=z), name="last")
         .build(record=True))
    return g, z


g, z = build(x.copy(), 0.8)
rev = eagle.diff.Reverse(g)                 # built once, used for both flag values

Launch with the flag on, reverse, flip the flag, launch and reverse again, and compare each answer with central differences of the same graph:

def J(x_, a_):
    g_, z_ = build(x_, a_)
    g_.launch()
    return z_.sum()


h = 1e-6
for on in (1, 0):
    flag[0] = on
    g.launch()
    grads = rev(g.record, {"last.y": np.ones(n)})
    fd_x = [(J(x + h * e, 0.8) - J(x - h * e, 0.8)) / (2 * h) for e in np.eye(n)]
    fd_a = (J(x, 0.8 + h) - J(x, 0.8 - h)) / (2 * h)
    print(f"flag={on}:", np.allclose(grads["x"], fd_x, rtol=1e-5, atol=1e-8),
          np.isclose(grads["a"], fd_a, rtol=1e-5, atol=1e-8), "dJ/da =", round(float(grads["a"]), 4))

flag=1: True True dJ/da = 0.0
flag=0: True True dJ/da = 2.1637

With the flag on, maybe overwrote y, so a had no effect and dJ/da is exactly 0. With the flag off, maybe was skipped, y stayed what first wrote, and the adjoint of y passed straight through the skipped step to first. That pass-through is the rule for a skipped step: a plane it would have updated keeps its earlier writer’s value, so its adjoint goes to that writer. eagle.repeat_while works the same way with an iteration count instead of a flag: the reverse unrolls exactly the iterations the launch ran.

A run too long for memory#

A recorded run-until-done (record=True) keeps the state every window steps and recomputes a window at a time in the reverse. memory_budget caps what that reverse holds, in bytes. Over the cap, eagle reverses the batch in passes of whole blocks of 256 samples (256, 512, 1024, …), and combines the passes with the same fixed tree a single pass would use. The gradients are therefore the same bits at any budget; only the time and the memory change.

@hawk.kernel(steps=1)
def pendulum(dt: Param, c: Param, t_final: Scalar, terminated: hawk.Terminated,
             x: Mutable[Scalar], v: Mutable[Scalar], t: Mutable[Scalar]):
    a = -m.sin(x) - c * v
    vn = v + dt * a
    x = x + dt * v
    v = vn
    tn = t + dt
    t = tn
    terminated = tn >= t_final


N = 2048


def gradients(budget=None):
    sim = eagle.simulation(pendulum, dt=0.01, c=0.3, t_final=np.linspace(0.5, 3.0, N),
                           x=np.linspace(0.5, 1.5, N), v=np.zeros(N), t=np.zeros(N),
                           max_steps=400, record=True, window=20, memory_budget=budget)
    sim.run()
    return eagle.diff.reverse(sim.record, {"x": np.ones(N)})


whole = gradients()                         # default: half the free memory
capped = gradients(budget=1_000_000)        # about 1 MB
print("passes:", whole.report.chunks, "->", capped.report.chunks)
print("same bits:", all(np.array_equal(whole[k], capped[k]) for k in whole))

passes: 1 -> 4
same bits: True

grads.report.chunks counts the passes. A budget too small for even one block of 256 samples is refused, and the message gives the bytes it needs and what to raise. Raising window or lowering max_steps also shrinks what one sample costs.

A stop between two steps#

An ordinary run-until-done stops a sample at the first step whose rule fires, so its stop time is a multiple of the step and moves by nothing when you nudge an input: the stop time’s derivative is zero. With events="refine" the run finds the crossing inside that last step and leaves the state at the event. The kernel names its time plane and step (time="t", time_step="dt"), and its stop rule is a smooth crossing, here y < 0 for a ball hitting the ground.

@hawk.kernel(steps=1, time="t", time_step="dt")
def fall(dt: Param, g: Param, terminated: hawk.Terminated,
         y: Mutable[Scalar], v: Mutable[Scalar], t: Mutable[Scalar]):
    yn = y + dt * v - 0.5 * g * dt * dt
    v = v - g * dt
    y = yn
    t = t + dt
    terminated = yn < 0.0


n, grav = 5, 9.81
y0 = np.linspace(1.0, 5.0, n)
t_hit = np.sqrt(2 * y0 / grav)              # the closed form


def drop(events):
    sim = eagle.simulation(fall, dt=0.1, g=grav, y=y0.copy(), v=np.zeros(n),
                           t=np.zeros(n), max_steps=200, events=events, record=True)
    sim.run()
    return sim


for events in ("step", "refine"):
    sim = drop(events)
    d_t = eagle.diff.reverse(sim.record, {"t": np.ones(n)})["y"]
    print(events, "dt*/dy0 =", d_t.round(4))
print("exact      ", (1 / (grav * t_hit)).round(4))
print("t* found   ", np.allclose(drop("refine").runner.t_stop, t_hit))

step dt*/dy0 = [0. 0. 0. 0. 0.]
refine dt*/dy0 = [0.2258 0.1596 0.1303 0.1129 0.101 ]
exact       [0.2258 0.1596 0.1303 0.1129 0.101 ]
t* found    True

The stepped run reports a zero sensitivity; the refined one matches the exact 1 / (g t*). runner.t_stop holds each sample’s event time and runner.event_status how it was located. A sample that never crosses inside its stop step is refused by name, unless you pass on_no_event="zero".

Hybrid trajectories: a Graph of phases#

A bouncing ball is a fall, a jump in velocity, then a rise. Each continuous piece is a refined run; eagle.diff.phase(run) makes one a node of an ordinary eagle.Graph, and the jump is a reset kernel between them. The reverse of the recorded graph chains the pieces, which is exactly the sensitivity of the jump (what the literature calls the saltation matrix). After one bounce with restitution e the apex is e² y0, so d(apex)/de = 2 e y0:

@hawk.kernel(steps=1, time="t", time_step="dt")
def rise(dt: Param, g: Param, terminated: hawk.Terminated,
         y: Mutable[Scalar], v: Mutable[Scalar], t: Mutable[Scalar]):
    vn = v - g * dt
    y = y + dt * v - 0.5 * g * dt * dt
    v = vn
    t = t + dt
    terminated = vn <= 0.0


@hawk.kernel
def bounce(e: Scalar, v_in: Scalar, v: Mutable[Scalar]):
    v = -e * v_in


(bounce_plan,) = eagle.deploy([bounce], targets=("host",), cache_dir=tempfile.mkdtemp())
e = np.full(n, 0.8)
y, t = y0.copy(), np.zeros(n)
v_down, v_up = np.zeros(n), np.zeros(n)
common = dict(dt=0.05, g=grav, y=y, t=t, max_steps=4000, events="refine", record=True)
down = eagle.simulation(fall, v=v_down, **common)
up = eagle.simulation(rise, v=v_up, **common)

graph = (eagle.Graph()
         .add_node(eagle.diff.phase(down), name="fall")
         .add(bounce_plan.bind(e=e, v_in=v_down, v=v_up), name="bounce")
         .add_node(eagle.diff.phase(up), name="rise")
         .build(record=True))
graph.launch()
grads = eagle.diff.reverse(graph.record, {"y": np.ones(n)})
print("apex = e^2 y0:", np.allclose(y, e**2 * y0))
print("d apex/de = 2 e y0:", np.allclose(grads["e"], 2 * e * y0))

apex = e^2 y0: True
d apex/de = 2 e y0: True

Both fall and rise share the planes y and t; each graph launch restarts every phase and runs it to its event. See Gradients through a graph: eagle.diff for what a refined reverse refuses and how the sensitivity is built.