Use eagle with JAX

Use eagle with JAX#

You write jax.jit(jax.grad(...)) steps, vmap over batches and lax.scan over time. This page changes one line of such a program: a step written in jax.numpy becomes a compiled hawk kernel. Your loop, your transformations and your numbers stay the same.

It runs on the CPU and on a CUDA GPU alike; the stream cell needs a CUDA build of JAX.

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 python -m raptor_docs.nb colab_setup (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}"])
    from packaging.version import Version

    try:
        old = Version(importlib.metadata.version("jax")) < Version("0.11.2")
    except importlib.metadata.PackageNotFoundError:
        old = True
    if old:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", f"jax{extra}>=0.11.2"])
    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

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 by gradient descent in one jitted training step. This is plain JAX; nothing from eagle yet.

import os

os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")  # leave GPU memory for your other libraries
try:  # CUDA 13: CuPy loads NVRTC before JAX does (see Limits)
    import cupy
    cupy.cuda.nvrtc.getVersion()
except Exception:  # no CuPy or no CUDA here
    pass
import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)  # the kernel runs in float64 (see Limits for float32)

N, STEPS, DT, G = 256, 20, 0.05, 9.81
v0 = 5.0 * jax.random.normal(jax.random.key(0), (3, N))


def jax_step(v, cd, dt, g):
    a = -cd * jnp.linalg.norm(v, axis=0) * v + jnp.array([0.0, 0.0, -g])[:, None]
    return v + dt * a


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


target = rollout(jax_step, 0.07)


def fit(step):
    loss = lambda cd: ((rollout(step, cd) - target) ** 2).mean()

    @jax.jit
    def train_step(cd):
        value, grad = jax.value_and_grad(loss)(cd)
        return cd - 2e-4 * grad, value

    cd, losses = 0.01, []
    for _ in range(60):
        cd, value = train_step(cd)
        losses.append(float(value))
    return float(cd), losses


cd_jax, losses_jax = fit(jax_step)
print(f"jax step:   cd = {cd_jax:.6f}, loss {losses_jax[0]:.4f} -> {losses_jax[-1]:.2e}")

jax step:   cd = 0.070000, loss 5.1412 -> 1.60e-15

2. Drop the kernel in#

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

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


@hawk.kernel
def drag_step(v: Vector[3], cd: Param, 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 jax as eagle_jax

step = eagle_jax.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]:.2e}")
assert abs(cd_eagle - cd_jax) < 1e-9 and abs(losses_eagle[-1] - losses_jax[-1]) < 1e-9
print("same numbers as the jax step")

eagle step: cd = 0.070000, loss 5.1412 -> 1.60e-15
same numbers as the jax step

3. What stays the same#

Your transformations still work on the changed program. grad under jit, vmap of grad over a few drag coefficients, and lax.scan over time, each against the jax.numpy step:

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


def scanned(s):
    return lambda cd: jax.lax.scan(lambda v, _: (s(v, cd, DT, G), None), v0, None, length=STEPS)[0]


cds = jnp.linspace(0.02, 0.1, 5)
for name, make in {
    "jit(grad)": lambda s: jax.jit(jax.grad(loss_of(s))),
    "vmap(grad)": lambda s: jax.vmap(jax.grad(loss_of(s))),
    "lax.scan": scanned,
}.items():
    arg = cds if name == "vmap(grad)" else 0.05
    assert jnp.allclose(make(step)(arg), make(jax_step)(arg), rtol=1e-9, atol=1e-12)
    print(f"{name:<11} matches the jax step")

jit(grad)   matches the jax step
vmap(grad)  matches the jax step
lax.scan    matches the jax step

Jacobians and second derivatives come from the same derived kernels, in either mode. Forward mode also runs through a kernel that carries state, and a kernel built with Guard(active_set=True) works in function() too (tutorial 13 runs both):

final = lambda s: (lambda cd: rollout(s, cd)[:, 0])  # the first ball's final velocity
for name, transform in {"jacfwd": jax.jacfwd, "jacrev": jax.jacrev}.items():
    assert jnp.allclose(transform(final(step))(0.05), transform(final(jax_step))(0.05), rtol=1e-9)
    print(f"{name:<8} d velocity / d cd =", jnp.round(transform(final(step))(0.05), 4))
h_eagle, h_jax = jax.hessian(loss_of(step))(0.05), jax.hessian(loss_of(jax_step))(0.05)
assert jnp.allclose(h_eagle, h_jax, rtol=1e-9)
print(f"hessian  d2 loss / d cd2 = {h_eagle:.4f} (jax step: {h_jax:.4f})")

jacfwd   d velocity / d cd = [ 5.5486 55.3816 64.3769]
jacrev   d velocity / d cd = [ 5.5486 55.3816 64.3769]
hessian  d2 loss / d cd2 = 2351.1435 (jax step: 2351.1435)

4. What crosses, and who owns it#

JAX buffers cross without a copy: the kernel reads your inputs in place and writes into a buffer XLA allocated. The result is an ordinary jax.Array, owned by JAX, and nothing you passed in is written to. A kernel that finishes samples returns the updated Terminated mask last; hand the old mask in with donate_argnums and the update happens in place.

from hawk import Terminated


@hawk.kernel(steps=1)
def land(r: Vector[3], terminated: Terminated, r_next: Mutable[Vector[3]]):
    r_next = r - vec(0.0, 0.0, 1.0)
    terminated = r_next[2] < 0.0


landing = eagle_jax.function(land, wrt=())
r = jnp.stack([jnp.zeros(8), jnp.zeros(8), jnp.arange(8.0) / 2])
done = jnp.zeros(8, dtype=bool)
r_in, done_in = jnp.array(r), jnp.array(done)

step_land = jax.jit(landing, donate_argnums=1)
mask_ptr = done.unsafe_buffer_pointer()
r_next, done_next = step_land(r, done)
print("result is a jax.Array:", isinstance(r_next, jax.Array), r_next.dtype, "| input untouched:", bool((r == r_in).all()))
print("mask updated in place:", done_next.unsafe_buffer_pointer() == mask_ptr, "| samples landed:", int(done_next.sum()))

result is a jax.Array: True float64 | input untouched: True
mask updated in place: True | samples landed: 2

On a GPU the launch rides XLA’s own stream, so it is ordered inside the jitted program and never waits for the host. Here a heavy matrix product runs first; the call still returns at once, and the kernel still reads its result:

import time

if jax.default_backend() == "gpu":
    a = jnp.ones((4096, 4096), dtype=jnp.float32)

    def heavy_then_step(a, v):
        for _ in range(6):
            a = a @ a / 4096.0
        return step(v * a[0, 0], 0.05, DT, G)  # a[0, 0] is 1.0 after the products

    run = jax.jit(heavy_then_step)
    run(a, v0).block_until_ready()  # compile once
    t0 = time.perf_counter()
    out = run(a, v0)
    t_dispatch = time.perf_counter() - t0
    out.block_until_ready()
    t_total = time.perf_counter() - t0
    assert t_dispatch < 0.2 * t_total and jnp.allclose(out, jax_step(v0, 0.05, DT, G))
    print("the call returned before the GPU finished; the result is still right")
else:
    print("this JAX has no CUDA device: the stream check runs on a CUDA build")

this JAX has no CUDA device: the stream check runs on a CUDA build

A total over all samples#

A kernel can also reduce over its samples: hawk.sum(x) stored to a Mutable is a cell, one value for the whole batch. The kernel writes one partial per 32 samples and the bridge finishes them into that value inside the same XLA call, so the result is a (1,) array that works under jit, grad and vmap like any other. jax.grad also works through a kernel that reads its total later (a split kernel).

function(kernel, policy=...) chooses how that total is finished: "compensated" (the default), "ordered" or "auto", as for eagle.reduce. Use "ordered" where the bits must be the same on every run; an unknown name is refused. Integer sums, max and min are exact under any policy.

from hawk import Scalar


@hawk.kernel
def energy(x: Scalar, total: Mutable[Scalar]):
    total = hawk.sum(x * x)


energy_of = eagle_jax.function(energy)
xs = jax.random.normal(jax.random.key(1), (1000,))
print("jit :", jax.jit(energy_of)(xs), " jnp:", (xs * xs).sum())
g = jax.jit(jax.grad(lambda a: energy_of(a).sum()))(xs)
assert jnp.allclose(g, 2 * xs)
print("grad matches 2x")

jit : [967.91464364]  jnp: 967.9146436427735
grad matches 2x

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] + '...'}")


refused("mixed dtypes", lambda: step(v0, jnp.ones(1, dtype=jnp.float32), DT, G))
refused("integer plane", lambda: step(v0.astype(int), 0.05, DT, G))
refused("wrong shape", lambda: step(v0[:2], 0.05, DT, G))
refused("missing input", lambda: step(v0, 0.05, DT))
refused("not an input", lambda: step(v0, 0.05, DT, G, drag=1.0))

integer plane: ModeError: eagle.frameworks.jax drag_step: plane 'v' is int64, which is not supported; a kernel runs float32 or float64 planes
wrong shape: ValueError: eagle.frameworks.jax: drag_step: input 'v' must have shape (3, 256) for n=256, got (2, 256)
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 JAX step inside an eagle batch#

The other direction: a jitted JAX function runs between eagle launches, over the same arrays. Here a JAX tanh feeds a hawk kernel whose weights live in one array that your optimizer updates:

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
gpu = jax.default_backend() == "gpu"
xp = __import__("cupy") if gpu else 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)
w = xp.array([2.0, 0.5])                           # JAX arrays are immutable: one array of your own holds the weights

to_np = (lambda a: a.get()) if gpu else np.asarray
model = eagle_jax.node(jax.jit(lambda x: jnp.tanh(0.7 * x)), reads={"x": x}, writes={"z": z})
g = (eagle.Graph().add_node(model, name="model")
     .add(plan.bind(w=w, z=z, y=y), name="lin").build(record=True))
w[:] = xp.from_dlpack(jnp.array([6.0, 1.5]))       # the optimizer's new weights, copied in: the next launch reads them
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

node hands your planes to JAX without a copy, copies JAX’s results into z on its own stream, and keeps JAX’s buffers alive until they have been used; the host never waits. On the CPU it is an eager call (the three tiers has the picture).

Your JAX code

Write it as

Runs

a jitted step

frameworks.jax.node(fn, reads=..., writes=...)

stream-ordered

a step that reads a value back

opaque(fn, eager=True)

host-synchronized

the whole batch inside jit

the bridge above (sections 1 to 5)

XLA’s own program

JAX is never captured into an eagle graph, because XLA owns its allocations; a stepping loop that runs a JAX step is therefore your own runner.run() with max_steps=k loop.

Weights. The Table is bound to one array of your own (w above); after each optimizer step copy the new value in.

Gradients. Build your vjp from jax.vjp of your step:

def vjp(primals, cotangents):
    _, pull = jax.vjp(lambda a: jnp.tanh(0.7 * a), jax.dlpack.from_dlpack(xp.asarray(primals["x"]).copy()))
    return {"x": xp.from_dlpack(pull(jax.dlpack.from_dlpack(xp.asarray(cotangents["z"]).copy()))[0])}


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 JAX step:", bool(np.allclose(to_np(grads["x"]), want)))

d y / d x through your JAX step: True

7. Limits#

What does not work today, one line each:

  • A call uses one GPU: the one its arrays live on (plans, graphs and reductions follow their planes; a JAX function runs on the JAX device it was built under, one function per GPU). Arrays on two GPUs in one call are refused by name.
  • A call runs the float32 or the float64 build, chosen by its first float plane, and mixing them is refused (the kernel is built once per dtype). JAX itself is float32 unless `jax_enable_x64` is on: cast every plane to one dtype, or turn x64 on as this page does.
  • `vmap` over a batch of `Param` values is one launch over every entry; a kernel that reduces over samples itself, or has a `jump=` plan, runs one launch per entry instead (our scope): put the varying value in a plane for one launch.
  • `jax.export` modules carry the kernel's digest, not its code (so they stay small and cache-friendly): in a fresh process call `eagle.frameworks.jax.preload(cache_dir)` before running one.
  • `Terminated` is a stop decision, not a number, so it carries no derivative and a finished sample's derivatives are zero. Keep any soft version of the decision in `jnp`, and differentiate the planes.
  • Needs `jax>=0.11.2`, the release whose XLA FFI version the bridge is built against (our scope); a newer jax works.
  • JAX cannot run inside a captured eagle graph or stepping loop, because XLA owns its allocations (XLA's design): use `eagle.frameworks.jax.node` between launches, or put the batch inside `jit`.
  • JAX leaves its last GPU current; eagle follows your arrays, so only NEW cupy allocations are affected (cupy allocates on the current device): make them inside `with cupy.cuda.Device(n):`.
  • On CUDA 13, CuPy cannot compile a kernel after JAX has started, because JAX loads NVIDIA's NVRTC library first and that build of it cannot find its own builtins (NVIDIA's and JAX's packaging): call `cupy.cuda.nvrtc.getVersion()` before your first JAX call, as the notebooks here do.
  • `jax.grad` does not differentiate through `lax.while_loop` (JAX's rule): use `lax.scan` with a next-state kernel.
  • JAX runs under free-threading on Python 3.14t only, because jaxlib publishes no 3.13t wheel (JAX's release choice): use 3.14t for free-threaded JAX, or a regular 3.13 build.

8. Where next#