eagle inside your JAX program

eagle inside your JAX program#

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

/dev/shm/jaxprobe/venv/lib/python3.12/site-packages/cupy/_environment.py:677: UserWarning: 
--------------------------------------------------------------------------------

  CuPy may not function correctly because multiple CuPy packages are installed
  in your environment:

    cupy, cupy-cuda12x

  Follow these steps to resolve this issue:

    1. For all packages listed above, run the following command to remove all
       existing CuPy installations:

         $ pip uninstall <package_name>

      If you previously installed CuPy via conda, also run the following:

         $ conda uninstall cupy

    2. Install the appropriate CuPy package.
       Refer to the Installation Guide for detailed instructions.

         https://docs.cupy.dev/en/stable/install.html

--------------------------------------------------------------------------------

  warnings.warn(f'''

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:

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")

the call returned before the GPU finished; the result is still right

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))

mixed dtypes: ModeError: eagle.frameworks.jax drag_step: mixed floating dtypes in one call (plane 'cd' is float32, the first plane is float64); pass every plane in one floa...
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:

  • One GPU per process until 0.7 (our scope): run one process per GPU, or choose it with `CUDA_VISIBLE_DEVICES`.
  • 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 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; over planes it is free for `in_axes=-1` and one transpose for a leading axis.
  • `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.
  • `function()` refuses a kernel built with `Guard(active_set=True)`, because its index map belongs to eagle's launcher (our scope): run that kernel with eagle, or build it without `active_set`.
  • 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`.
  • Starting JAX makes its last GPU the thread's current device, and each JAX call then makes its own GPU current (JAX's behavior); eagle graphs and `jax.node` still run on their planes' GPU: make other cupy arrays inside `with cupy.cuda.Device(n):`, or show one GPU with `CUDA_VISIBLE_DEVICES`.
  • 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#