eagle inside your PyTorch program

eagle inside your PyTorch program#

You train with torch.optim, call loss.backward(), and keep your tensors on whatever device you use. This page changes one line of such a program: a step written in torch becomes a compiled kernel. The loop, the optimizer and the numbers stay the same, and so do torch.compile, torch.func and CUDA graphs.

It runs on the CPU and on a CUDA GPU alike; the graph-capture cell needs a CUDA build of torch.

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

1. Your program as it is#

A ball-flight step with air drag, rolled out 20 steps for 256 balls, and a drag coefficient fitted with Adam. This is plain PyTorch; nothing from eagle yet.

import warnings

warnings.filterwarnings("ignore", category=FutureWarning, message=r"The pynvml package is deprecated.*")
warnings.filterwarnings("ignore", category=UserWarning, message=r"(?s).*CuPy may not function correctly.*")
import torch

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"
N, STEPS, DT, G = 256, 20, 0.05, 9.81
v0 = 5.0 * torch.randn(3, N, dtype=torch.float64, device=device)


def torch_step(v, cd, dt, g):
    a = -cd * v.norm(dim=0) * v
    a = a + torch.tensor([[0.0], [0.0], [-g]], dtype=torch.float64, device=v.device)
    return v + dt * a


def rollout(step, cd):
    v = v0
    for _ in range(STEPS):
        v = step(v, cd.expand(N), DT, G)
    return v


target = rollout(torch_step, torch.tensor(0.07, dtype=torch.float64, device=device)).detach()


def fit(step):
    cd = torch.tensor(0.01, dtype=torch.float64, device=device, requires_grad=True)
    opt = torch.optim.Adam([cd], lr=0.01)
    losses = []
    for _ in range(60):
        opt.zero_grad()
        loss = ((rollout(step, cd) - target) ** 2).mean()
        loss.backward()
        opt.step()
        losses.append(loss.item())
    return cd.item(), losses


cd_torch, losses_torch = fit(torch_step)
print(f"torch step: cd = {cd_torch:.6f}, loss {losses_torch[0]:.4f} -> {losses_torch[-1]:.6f}")

torch step: cd = 0.070443, loss 4.7666 -> 0.000005

2. Drop the kernel in#

The same step as a hawk kernel, and the one line that turns it into a torch function. fit and rollout are not touched.

import hawk
from hawk import Mutable, Param, Scalar, Vector
from hawk.math import norm, vec


@hawk.kernel
def drag_step(v: Vector[3], cd: Scalar, dt: Param, g: Param, v_next: Mutable[Vector[3]]):
    a = vec(0.0, 0.0, -g) - cd * norm(v) * v
    v_next = v + dt * a

from eagle.frameworks import torch as eagle_torch

step = eagle_torch.function(drag_step, wrt=("v", "cd"))  # the one changed line

cd_eagle, losses_eagle = fit(step)
print(f"eagle step: cd = {cd_eagle:.6f}, loss {losses_eagle[0]:.4f} -> {losses_eagle[-1]:.6f}")
assert abs(cd_eagle - cd_torch) < 1e-9 and abs(losses_eagle[-1] - losses_torch[-1]) < 1e-9
print("same numbers as the torch step")

eagle step: cd = 0.070443, loss 4.7666 -> 0.000005
same numbers as the torch step

3. What stays the same#

Your own tools still work on the changed program. backward gives the same gradient as autograd through the torch step:

def grad_of(step):
    cd = torch.tensor(0.05, dtype=torch.float64, device=device, requires_grad=True)
    ((rollout(step, cd) - target) ** 2).mean().backward()
    return cd.grad


g_torch, g_eagle = grad_of(torch_step), grad_of(step)
torch.testing.assert_close(g_eagle, g_torch, rtol=1e-9, atol=1e-12)
print("d loss / d cd:", g_eagle.item(), "(torch step:", g_torch.item(), ")")

d loss / d cd: -30.672293533875695 (torch step: -30.672293533875695 )

Forward-mode AD works too, through torch.autograd.forward_ad:

import torch.autograd.forward_ad as fwAD

cd0 = torch.full((N,), 0.05, dtype=torch.float64, device=device)
with fwAD.dual_level():
    dual = fwAD.make_dual(cd0, torch.ones_like(cd0))
    tangent = fwAD.unpack_dual(step(v0, dual, DT, G)).tangent
eps = 1e-6
fd = (step(v0, cd0 + eps, DT, G) - step(v0, cd0 - eps, DT, G)) / (2 * eps)
torch.testing.assert_close(tangent, fd, rtol=1e-5, atol=1e-8)
print("forward-mode tangent matches a finite difference, shape", tuple(tangent.shape))

forward-mode tangent matches a finite difference, shape (3, 256)

torch.compile traces straight through the step, forward and backward, with no graph break. (This page uses the aot_eager backend so it also runs on older GPUs; on a recent one use the default.)

import torch._dynamo as dynamo


def loss_of(s):
    return lambda cd: ((rollout(s, cd) - target) ** 2).mean()


cd = torch.tensor(0.05, dtype=torch.float64, device=device, requires_grad=True)
print("graph breaks:", dynamo.explain(loss_of(step))(cd).graph_break_count)

compiled = torch.compile(loss_of(step), fullgraph=True, backend="aot_eager")
compiled(cd).backward()
torch.testing.assert_close(cd.grad, g_torch, rtol=1e-9, atol=1e-12)
print("compiled gradient matches the torch step:", cd.grad.item())

graph breaks: 0
compiled gradient matches the torch step: -30.672293533875695

Derivatives of derivatives, torch.func and float32 work as they do on any torch function:

cd = torch.tensor(0.05, dtype=torch.float64, device=device, requires_grad=True)
h_eagle = torch.autograd.functional.hessian(loss_of(step), cd)
h_torch = torch.autograd.functional.hessian(loss_of(torch_step), cd)
torch.testing.assert_close(h_eagle, h_torch, rtol=1e-8, atol=0)
print("hessian:", h_eagle.item(), "(torch step:", h_torch.item(), ")")

cds = torch.linspace(0.02, 0.1, 5, dtype=torch.float64, device=device)
per_cd = torch.func.vmap(torch.func.grad(loss_of(step)))(cds)
torch.testing.assert_close(per_cd, torch.func.vmap(torch.func.grad(loss_of(torch_step)))(cds), rtol=1e-8, atol=1e-12)
print("torch.func.vmap(grad) over 5 drag coefficients matches the torch step")

out32 = step(v0.float(), torch.full((N,), 0.05, dtype=torch.float32, device=device), DT, G)
print("float32 in, float32 out:", out32.dtype)

hessian: 2179.2612438437027 (torch step: 2179.2612438437022 )
torch.func.vmap(grad) over 5 drag coefficients matches the torch step
float32 in, float32 out: torch.float32

4. What crosses, and who owns it#

Tensors cross without a copy, in both directions. The result is a new torch tensor, owned by torch, on the device of its inputs. Nothing you passed in is written to, and a Param that is a tensor is read from the device, never copied to the host.

v_in = v0.clone()
out = step(v0, torch.full((N,), 0.05, dtype=torch.float64, device=device), DT, G)
print("result:", type(out).__name__, out.dtype, out.device.type, "| a new tensor:", out.data_ptr() != v0.data_ptr())
print("input untouched:", bool((v0 == v_in).all()))

if torch.cuda.is_available():
    # the launch rides torch's current stream: queued behind a sleep, it still reads the data written after it
    s = torch.cuda.Stream()
    cd_gpu = torch.full((N,), 0.05, dtype=torch.float64, device="cuda")
    with torch.cuda.stream(s):
        torch.cuda._sleep(100_000_000)
        t = torch.zeros(3, N, dtype=torch.float64, device="cuda")
        t.fill_(2.0)
        y = step(t, cd_gpu, DT, G)
    torch.cuda.synchronize()
    torch.testing.assert_close(y, torch_step(t, cd_gpu, DT, G))
    print("on a side stream, no sync: the launch ran after the fill, in order")
else:
    print("this torch has no CUDA: the stream check runs on a CUDA build")

result: Tensor torch.float64 cuda | a new tensor: True
input untouched: True
on a side stream, no sync: the launch ran after the fill, in order

On a GPU the whole gradient step can be captured in a CUDA graph and replayed; a changed Param is a plain write into the same tensor.

def grad_of_at(value):
    c = torch.tensor(value, dtype=torch.float64, device=device, requires_grad=True)
    loss_of(torch_step)(c).backward()
    return c.grad


if torch.cuda.is_available():
    cd = torch.tensor(0.05, dtype=torch.float64, device="cuda", requires_grad=True)
    loss = loss_of(step)
    loss(cd).backward()  # warm up: build and load the kernels
    cd.grad = None
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        loss(cd).backward()
    for value in (0.05, 0.08):
        with torch.no_grad():
            cd.fill_(value)
        graph.replay()
        torch.cuda.synchronize()
        torch.testing.assert_close(cd.grad, grad_of_at(value), rtol=1e-9, atol=1e-12)
    print("captured once, replayed with two different drag coefficients: gradients match")
else:
    print("this torch has no CUDA: the capture check runs on a CUDA build")

captured once, replayed with two different drag coefficients: gradients match

5. What is refused, and why#

A kernel reads and writes exact buffers, so anything it cannot use as it is gets a named refusal.

def refused(label, call):
    try:
        call()
    except Exception as err:
        msg = str(err).splitlines()[0]
        print(f"{label}: {type(err).__name__}: {msg if len(msg) <= 150 else msg[:147] + '...'}")


cd_ok = torch.full((N,), 0.05, dtype=torch.float64, device=device)
refused("mixed dtypes", lambda: step(v0, cd_ok.float(), DT, G))
refused("wrong length", lambda: step(v0, cd_ok[:10], DT, G))
refused("missing input", lambda: step(v0, cd_ok, DT))
refused("not an input", lambda: step(v0, cd_ok, DT, G, drag=1.0))

mixed dtypes: ValueError: eagle.frameworks.torch: mixed floating dtypes in one call (torch.float32, torch.float64); pass every plane and Param as float64, or every one as fl...
wrong length: ValueError: drag_step: input 'cd' must have shape (256,) for n=256, got (10,)
missing input: TypeError: drag_step: missing input(s) ['g']
not an input: TypeError: drag_step: 'drag' is not an input (inputs: ['v', 'cd', 'dt', 'g'])

6. Your torch model inside an eagle batch#

The other direction: your torch code runs as a node of an eagle graph, in the middle of a kernel pipeline, over the same arrays. Here a torch tanh feeds a hawk kernel whose two weights are your parameter:

import tempfile
import warnings

warnings.filterwarnings("ignore", category=UserWarning, message=r"(?s).*CuPy may not function correctly.*")

import numpy as np

import eagle, eagle.diff, eagle.plan
import hawk
from eagle.graph import opaque
from eagle.interop import import_buffer
from hawk import Mutable, Scalar, Table
import torch

gpu = torch.cuda.is_available()
dev, xp = ("cuda", __import__("cupy")) if gpu else ("cpu", np)
@hawk.kernel
def lin(w: Table[Scalar], z: Scalar, y: Mutable[Scalar]):
    y = w.at(0) * z + w.at(1)       # w is ONE buffer, read by every sample

(plan,) = eagle.plan.auto([lin], targets=("cuda" if gpu else "host",), cache_dir=tempfile.mkdtemp())
x, z, y = xp.linspace(-1, 1, 8), xp.zeros(8), xp.zeros(8)
c = torch.tensor(0.7, dtype=torch.float64, device=dev, requires_grad=True)           # your model's parameter
w = torch.nn.Parameter(torch.tensor([2.0, 0.5], dtype=torch.float64, device=dev))    # the kernel's weights
tx, tz = torch.as_tensor(x, device=dev), torch.as_tensor(z, device=dev)              # views of the planes, no copy


def model():
    with torch.no_grad():
        tz.copy_(torch.tanh(c * tx))                                                 # your torch code

to_np = (lambda a: a.get()) if gpu else np.asarray
g = (eagle.Graph()
     .add_node(opaque(model, reads={"x": x}, writes={"z": z}, eager=gpu), name="model")
     .add(plan.bind(w=import_buffer(w), z=z, y=y), name="lin").build(record=True))
with torch.no_grad():
    w.mul_(3.0)                       # the optimizer's in-place step is the refresh: no copy, no rebuild
g.launch()
print("a launch after the build sees the new weights:", bool(np.allclose(to_np(y), 6 * np.tanh(0.7 * to_np(x)) + 1.5)))

a launch after the build sees the new weights: True

The graph worked out the order from the planes you declared (x into model, z into lin), so there is no stream or event to manage. How much eagle can do for your code depends on one keyword; for torch it reads like this (the three tiers has the picture):

Your torch code	Write it as	Runs
an op with `out=` or in place, no allocation	`opaque(fn)`	captured with the graph
a whole module, captured with `torch.cuda.CUDAGraph(keep_graph=True)`	`child(graph, reads=..., writes=...)`	captured, weights live
any other torch code	`opaque(fn, stream=torch.cuda.Stream())`	stream-ordered, the host never waits
code that reads a value back (`.item()`)	`opaque(fn, eager=True)`	host-synchronized: the GPU idles while it runs

The cell above uses eager=True on the GPU because tanh allocates; on the CPU every tier is a plain call. In eagle.until_done([model, plan], ...) the first two rows run once per step, over the whole batch; the last two cannot, because a stepping loop is one device graph, so call runner.run() with max_steps=k, then your code, in your own loop.

Weights are one buffer. import_buffer(param) binds the parameter itself to the Table, so optimizer.step() is the refresh. Update weights after launch() returns (or synchronize()), never between enqueue and synchronize; a launch after param.data = ... is refused by name, so bind the tensor your optimizer updates.

Gradients through your model. Name the node and hand reverse your vjp under that name. It receives the recorded inputs and the incoming cotangents and returns one gradient per read; your own parameters accumulate their .grad as you write them:

def vjp(primals, cotangents):
    x_ = torch.as_tensor(xp.asarray(primals["x"]).copy(), device=dev)
    ct = torch.as_tensor(xp.asarray(cotangents["z"]).copy(), device=dev)
    gx, gc = torch.autograd.grad(torch.tanh(c * x_.requires_grad_(True)), [x_, c], ct)
    c.grad = gc if c.grad is None else c.grad + gc                    # accumulate across reverses
    return {"x": xp.asarray(gx)}


grads = eagle.diff.reverse(g.record, {"lin.y": xp.ones(8)}, vjps={"model": vjp})
want = 6 * 0.7 * (1 - np.tanh(0.7 * to_np(x)) ** 2)
print("d y / d x through your model:", bool(np.allclose(to_np(grads["x"]), want)), "| c.grad set:", c.grad is not None)

d y / d x through your model: True | c.grad set: True

A reverse costs what your model costs per step, and the node must be pure (model.eval(), no random draw), because the reverse re-runs it: see your model inside a recorded loop, which also takes the same vjps= for a whole stepping loop recorded with until_done(..., record=True).

7. Limits#

What does not work today, one line each:

One GPU per process until 0.7 (our scope): run one process per GPU, or choose it with `CUDA_VISIBLE_DEVICES`.
`torch.compile` needs torch 2.4 or newer, the release that added the custom-op API the bridge uses; on an older torch run the step eagerly.
Inside a `torch.compile` region only first-order reverse mode runs, because compile traces the backward once, ahead of time (torch's design); run forward AD and double backward outside it, or through `torch.func`.
A kernel that finishes its own samples returns a fresh `Terminated` mask and leaves yours untouched (our design, so inputs are never written); keep the returned mask as the next step's.
`torch.func.vmap` over a batch of `Param` values runs one launch per entry, because a `Param` is one number per launch (our scope); put the varying value in a plane for one launch.
In a recorded loop a node runs before the plan, preparing the step from the current state (our scope until 0.6.3): write the step as `[node, plan]`.
In a recorded loop a node cannot read a plane another node writes (our scope until 0.6.3): fold the two into one node.
In a recorded loop a node reads per-sample planes, not a `Param` (our scope until 0.6.3): bind the value as a per-sample plane or a `Table`.
A recorded loop with a node stops with `events="step"` (our scope until 0.6.3): `events="refine"` re-steps the last step without your node.
In a recorded loop the node's vjp is a callable run host-synchronized once per step (our scope until 0.6.3): a captured backward as a `Node` vjp is refused by name.
A torch module that allocates cannot be captured as a plain `opaque`, because torch hands its graph memory to the next replay (torch's design): capture it with `torch.cuda.CUDAGraph(keep_graph=True)` and insert it with `child`, or write it with `out=`.
A stream or host node cannot run inside a stepping loop's device graph (our scope): capture it (`opaque(fn)`, or `child`), or call `runner.run()` with `max_steps=k`, then your code, in your own loop.
An eagle graph cannot be captured inside `torch.cuda.graph` or `wp.capture_begin` (our scope): build the eagle graph first and replay it next to your framework's.
torch's CPU wheels cover free-threaded Python 3.14t only, none for 3.13t (torch's release choice): use 3.14t for free-threaded torch, or a regular 3.13 build.

8. Where next#

Gradients through a torch loop: the same bridge, one kernel, step by step.
Interoperability: every framework, and the certification table that says what is checked end to end.
CuPy and NumPy and NVIDIA Warp: the same one-line change in those workflows.