Gradients through a graph: eagle.diff

Gradients through a graph: `eagle.diff`#

Run a graph forward once, then ask how its outputs move with every input, for every sample, in one backward sweep.

Time: ~10 min · Runs on: CPU (the same cells capture one CUDA graph for the reverse when the planes are cupy arrays) · You need: eagle, hawk and NumPy; read first: Build a step as a graph: eagle.Graph, Gradients through a graph: eagle.diff

You will end with gradients checked against finite differences, a fork whose two branches add up correctly, and a small optimisation loop that reuses one reverse pass.

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

import tempfile

import numpy as np
import eagle

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

A graph to differentiate#

Two small hawk kernels: y = a x², then z = sin(y). a is a uniform: one number every sample shares. The only new thing compared with Build a step as a graph: eagle.Graph is build(record=True): it tells the graph to keep what a reverse pass will need.

@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)
y, z = np.zeros(n), np.zeros(n)

g = (eagle.Graph()
     .add(square_plan.bind(x=x, a=0.5, y=y), name="square")
     .add(wave_plan.bind(x=y, y=z), name="wave")
     .build(record=True))
g.launch()
print(z)
print(g.record)

[0.47942554 0.17902957 0.01999867 0.01999867 0.17902957 0.47942554]
<eagle.diff.Record generation 1, live>

g.record is the Record of that launch: what its reverse reads. Its generation is 1, because it is the first launch.

The reverse pass#

Say our objective is the plain sum of the outputs, J = Σ zᵢ. Reverse mode (reverse mode) needs one thing from us: how J moves with each output, ∂J/∂zᵢ. That is the seed, one value per sample, and for a plain sum every seed is 1. (JAX calls it a cotangent, PyTorch grad_outputs, CasADi adj_seed.)

Both nodes call their output y, so the output of wave is named "wave.y".

grads = eagle.diff.reverse(g.record, {"wave.y": np.ones(n)})
print(list(grads))
print("dJ/dx =", grads["x"])
print("dJ/da =", grads["a"])
print(grads.report)

['x', 'a']
dJ/dx = [-0.87758256 -0.59030622 -0.19996     0.19996     0.59030622  0.87758256]
dJ/da = 2.5435165831215194
Report(generation=1, bitwise={'x': True, 'a': True})

Two gradients came back. grads["x"] has one value per sample, in sample order, because x is a per-sample input. grads["a"] is one number: a is shared, so each sample’s contribution is summed. Nobody wrote a derivative: eagle built each kernel’s vjp with hawk.diff.vjp and ran them backwards, wave first, then square.

Is it right?#

z = sin(a x²), so ∂z/∂x = cos(a x²) · 2 a x and ∂J/∂a = Σ cos(a x²) x². Check both, and also against central differences, which need nothing but the forward:

a = 0.5
print(np.allclose(grads["x"], np.cos(a * x * x) * 2 * a * x, rtol=1e-12))
print(np.isclose(grads["a"], np.sum(np.cos(a * x * x) * x * x), rtol=1e-12))


def forward(x_, a_):
    y_, z_ = np.zeros(n), np.zeros(n)
    (eagle.Graph().add(square_plan.bind(x=x_, a=a_, y=y_))
     .add(wave_plan.bind(x=y_, y=z_)).build().launch())
    return z_


h = 1e-6
fd_x = (forward(x + h, a) - forward(x - h, a)) / (2 * h)
fd_a = np.sum(forward(x, a + h) - forward(x, a - h)) / (2 * h)
print(np.allclose(grads["x"], fd_x, rtol=1e-6), np.isclose(grads["a"], fd_a, rtol=1e-6))

True
True
True True

Two branches reading one input#

In a fork, two nodes read the same plane. Its adjoint (how J moves with it) is then the sum of what flows back from both branches. eagle adds them in one fixed order, the order you added the branches, so the answer is the same bits on every run.

@hawk.kernel
def join(u: Scalar, v: Scalar, w: Mutable[Scalar]):
    w = u * v + u


(join_plan,) = eagle.deploy([join], targets=("host",), cache_dir=tempfile.mkdtemp())
u, v, w = np.zeros(n), np.zeros(n), np.zeros(n)
fork = (eagle.Graph()
        .add_group([square_plan.bind(x=x, a=a, y=u), wave_plan.bind(x=x, y=v)],
                   names=["square", "wave"])
        .add(join_plan.bind(u=u, v=v, w=w), name="join")
        .build(record=True))
fork.launch()
fork.show()

gf = eagle.diff.reverse(fork.record, {"w": np.ones(n)})
want = (np.sin(x) + 1) * 2 * a * x + a * x * x * np.cos(x)   # both branches
print(np.allclose(gf["x"], want, rtol=1e-12))

eagle.Graph  3 nodes · 1 group · backend: cpu · recorded

  stage 0 branch_tasks ┌─────────┐   ┌───────┐
                       │ square  │   │ wave  │
                       └────┬────┘   └───┬───┘
                            y            y
                            └─────┬──────┘
  stage 1 work_share          ┌───┴───┐
                              │ join  │
                              └───────┘
True

One Record per launch#

A Record points at the graph’s own planes, which the next launch overwrites. So the graph will not record a new launch while the last Record is still waiting for its reverse:

g2 = (eagle.Graph()
      .add(square_plan.bind(x=x.copy(), a=a, y=np.zeros(n)), name="square")
      .build(record=True))
g2.launch()
first = g2.record
try:
    g2.launch()
except RuntimeError as err:
    print(err)

kept = first.keep()        # a private copy: the graph is free to run again
g2.launch()
print(first.generation, kept.generation, g2.record.generation)
print(eagle.diff.reverse(kept, {"y": np.ones(n)})["x"])

eagle.Graph.launch: the Record of the previous launch is still live; reverse it, or record.keep() to keep a private copy
1 1 2
[-1.  -0.6 -0.2  0.2  0.6  1. ]

Reversing first now would be refused: its launch is no longer the newest, and eagle will not hand back a gradient computed from a mix of two launches. kept stays valid for good.

Built once, run in a loop#

eagle.diff.reverse builds a reverse per call. In a loop, build an eagle.diff.Reverse once and call it; on the GPU it is captured once and replayed. Here we move each x so that z = sin(a x²) reaches a target, by plain gradient descent on J = ½ Σ (zᵢ - targetᵢ)², whose seed is z - target. wrt=["x"] (wrt) asks for just the gradient we use.

target = np.full(n, 0.3)
rev = eagle.diff.Reverse(g)
for step in range(200):
    g.launch()
    grads = rev(g.record, {"wave.y": z - target}, wrt=["x"])
    x -= 0.5 * grads["x"]       # the graph reads x where it lies: no rebuild
print(np.round(z, 6))

[0.3 0.3 0.3 0.3 0.3 0.3]

An optimisation loop over a parameter#

Now tune the shared parameter a instead, so that the average of z hits 0.2: J = ½ (mean(z) - 0.2)², whose seed is (mean(z) - 0.2) / n for every sample. g.update(a=...) changes a between launches. A recorded graph keeps each parameter in a one-cell plane that the kernels and their vjps read, so this is a copy into that cell: nothing is rebuilt, and on the GPU nothing is recaptured, forward or reverse.

a_fit = 0.5
for step in range(100):
    g.update(a=a_fit)
    g.launch()
    mean = z.mean()
    grads = rev(g.record, {"wave.y": np.full(n, (mean - 0.2) / n)}, wrt=["a"])
    a_fit -= 2.0 * float(grads["a"])
print(round(a_fit, 6), round(z.mean(), 6))

0.330428 0.2

Same bits, every time#

Every sum eagle does on the way back has one fixed order, and a parameter’s gradient is summed in sample order, so a gradient is bitwise reproducible: the same bits for every thread count and every run. grads.report.bitwise says so per plane; the only gradients that can say False come from a kernel that reads a table by index, and strict=True refuses those instead.

fork.launch()                     # x has moved since: record a fresh launch
r1 = eagle.diff.reverse(fork.record, {"w": np.ones(n)})
r2 = eagle.diff.reverse(fork.record, {"w": np.ones(n)})
print(r1["x"].tobytes() == r2["x"].tobytes(), r1.report.bitwise)

True {'x': True, 'a': True}

What just happened#

build(record=True) made the graph keep what its reverse needs: your own planes, plus a copy of any plane a later node overwrites.
eagle.diff.reverse(g.record, seeds) ran each node’s vjp backwards and returned the gradient with respect to every input: per sample for per-sample inputs, summed for shared parameters.
In a fork, the adjoints of the shared input were summed in a fixed order.
A Record names one launch; keep() keeps a copy when you need it longer.
eagle.diff.Reverse is the built-once form for loops.

Try this#

Seed "wave.y" with 2 * z: that is the gradient of J = Σ zᵢ².
Pass seeds={"x": ...} and read the message: it tells you that ∂J/∂x is read from the result, not seeded.
Move the planes to the GPU (cupy arrays and targets=("cuda",)): the same cells capture the reverse into one CUDA graph.