A JAX model between eagle launches#
Your JAX model decides, eagle steps 512 oscillators, and the two agree to rounding on the CPU and on a GPU.
node = eagle_jax.node(model, reads={"x": x, "v": v}, writes={"corr": corr}) # your JAX model, as a graph node
graph = eagle.Graph().add_node(node, name="model").add(osc.bind(...), name="osc").build()
graph.launch(200) # model, step, model, step, ...
Time: ~5 min · Runs on: CPU, GPU if visible (picked up automatically) · You need: eagle, hawk and 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.
import sys, pathlib
sys.path.insert(0, str(pathlib.Path.cwd().parent.parent / "_shared"))
from nb_helpers import plot_style
plot_style()
Your model, their step#
The model is a small network written in jax.numpy. It looks at each oscillator’s position and velocity and returns
a correction. The step that moves the oscillators is a hawk kernel that takes that correction as an input plane.
Nothing here is specific to eagle yet: model is an ordinary jitted JAX function.
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 on the JAX page)
import cupy
cupy.cuda.nvrtc.getVersion()
except Exception: # no CuPy or no CUDA here
pass
import jax
import jax.numpy as jnp
import numpy as np
jax.config.update("jax_enable_x64", True)
import hawk
from hawk import Mutable, Param, Scalar, Terminated
import eagle
import eagle.plan
from eagle.frameworks import jax as eagle_jax
N, DT, C = 512, 0.01, 0.3
k1, k2 = jax.random.split(jax.random.key(0))
W1, W2 = 0.5 * jax.random.normal(k1, (2, 8)), 0.5 * jax.random.normal(k2, (8,))
@jax.jit
def model(x, v):
return jnp.tanh(jnp.stack([x, v], 1) @ W1) @ W2 # (N,) correction per oscillator
@hawk.kernel(steps=1)
def osc(dt: Param, c: Param, t_final: Scalar, corr: Scalar, terminated: Terminated,
x: Mutable[Scalar], v: Mutable[Scalar], t: Mutable[Scalar]):
vn = v + dt * (-x - c * v + corr * v)
x = x + dt * v
v = vn
tn = t + dt
t = tn
terminated = tn >= t_final # each oscillator stops at its own final time
One graph, two owners#
Where the planes live decides where the step runs: NumPy arrays for the CPU, CuPy arrays for a GPU. The same lines build the graph either way, with the JAX node first and the kernel after it, so each launch is one model call and one step.
gpu = jax.default_backend() == "gpu"
if gpu:
import cupy as xp
else:
xp = np
(plan,) = eagle.plan.auto([osc], targets=("cuda" if gpu else "host",))
start = dict(x=np.linspace(0.5, 1.5, N), v=np.linspace(-0.2, 0.3, N), t=np.zeros(N),
corr=np.zeros(N), t_final=np.linspace(0.5, 2.0, N))
planes = {name: xp.asarray(a).copy() for name, a in start.items()}
mask = xp.zeros(N, bool)
node = eagle_jax.node(model, reads={"x": planes["x"], "v": planes["v"]}, writes={"corr": planes["corr"]})
graph = (eagle.Graph()
.add_node(node, name="model")
.add(plan.bind(dt=DT, c=C, terminated=mask, finished_count=xp.zeros(1, np.uint32), **planes), name="osc")
.build())
graph.schedule
('eager_stream', 'capture')
schedule says how each piece runs. On a GPU the middle entry reads eager_stream: JAX runs on a stream of its own
(a queue of GPU work), eagle orders it between the two kernel segments on the device, and the host never waits. On the
CPU there are no streams, so the whole graph is one host schedule.
Run it, and check it against plain JAX#
We launch one step at a time and keep five oscillators’ positions to draw. The reference is the same loop written
entirely in jax.numpy, with a where that freezes each oscillator once it has stopped.
def host(a):
return a.get() if hasattr(a, "get") else np.asarray(a)
track = [host(planes["x"])[::128].copy()]
for _ in range(200): # 200 steps of 0.01 reach the longest final time, 2.0
graph.launch()
track.append(host(planes["x"])[::128].copy())
graph.synchronize()
@jax.jit
def reference_step(x, v, t, done):
vn = v + DT * (-x - C * v + model(x, v) * v)
xn, tn = x + DT * v, t + DT
return (jnp.where(done, x, xn), jnp.where(done, v, vn), jnp.where(done, t, tn),
done | (tn >= jnp.asarray(start["t_final"])))
state = (jnp.asarray(start["x"]), jnp.asarray(start["v"]), jnp.zeros(N), jnp.zeros(N, bool))
ref_track = [np.asarray(state[0])[::128]]
for _ in range(200):
state = reference_step(*state)
ref_track.append(np.asarray(state[0])[::128])
err = float(np.abs(host(planes["x"]) - np.asarray(state[0])).max())
print(f"all stopped: {bool(host(mask).all())}; largest difference from plain JAX: {err:.1e}")
assert err < 1e-12
all stopped: True; largest difference from plain JAX: 6.7e-16
import matplotlib.pyplot as plt
time = np.arange(201) * DT
fig, ax = plt.subplots(figsize=(6.5, 3.4))
ax.plot(time, np.array(track), lw=1.8)
ax.plot(time, np.array(ref_track), "k--", lw=0.8, label="plain JAX")
ax.set(xlabel="time", ylabel="position", title="Five oscillators, each frozen at its own final time")
ax.legend(loc="upper right")
plt.tight_layout()
The solid lines are eagle stepping with your model in the loop, the dashed ones the all-JAX loop: they sit on top of each other, and each curve goes flat where its oscillator stopped.
Why the results stay correct#
JAX hands the model’s result to eagle without copying. eagle then copies it into corr on the node’s own stream, and
holds on to JAX’s buffer until that copy has really run, so JAX may reuse its memory for the next call without racing
the copy. You write nothing for this; it is the reason node exists instead of a bare callback. Jit the model, as
above: an un-jitted call spends host time on every launch.
Inside until_done#
On the CPU the same node and kernel also run inside until_done, which stops each sample as it finishes. On a GPU
that loop is one device graph, so a JAX node is refused with a message that says so.
try:
again = {name: xp.asarray(a).copy() for name, a in start.items()}
inner = eagle_jax.node(model, reads={"x": again["x"], "v": again["v"]}, writes={"corr": again["corr"]})
report = eagle.until_done([inner, plan], max_steps=400, dt=DT, c=C, terminated=xp.zeros(N, bool), **again).run()
print("until_done:", report.steps, "steps,", "done" if report.done else "not done")
except ValueError as refusal:
print(str(refusal).split(": run host")[0])
eagle.until_done: node 'opaque_stream' runs on the host each launch, but a stepping loop is one device graph
When a JAX call has to read values back on the host, use opaque(fn, eager=True) instead: it works anywhere, at the
price of waiting for the device twice per launch. Your own code in the graph lays out all three tiers.
Limits#
- On a GPU a JAX node runs between launches, not inside `until_done`'s one device graph, because XLA owns its own memory and cannot be captured (XLA's design): drive the graph with your own loop, as this page does, or put the batch inside `jit` (tutorial 18).
Next#
The eagle batch inside
jit: the other direction, when JAX owns the loop.eagle inside your JAX program: what crosses, what is refused, the limits.