A stepping loop in lax.scan#
256 balls fly until they land. The kernel decides who has landed and says so in a mask; lax.scan carries the mask
to the next step, so a landed ball is never touched again.
Time: ~5 min · Runs on: CPU, GPU if visible (picked up automatically) · You need: eagle, hawk and JAX
You will build: a scan that stops each ball where it lands, the same loop on the CPU and the GPU with identical numbers, and a gradient through the whole loop.
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.
/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'''
A kernel that finishes its own samples#
terminated = r_next[2] < 0.0 marks a ball as landed. The bridge returns that mask as the last output; hand it back
in at the next step and the kernel skips every finished ball.
import os
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") # leave GPU memory for your other libraries
import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True) # the kernel runs in float64
import hawk
from hawk import Mutable, Param, Terminated, Vector
from hawk.math import norm, vec
from eagle.frameworks import jax as eagle_jax
@hawk.kernel(steps=1)
def flight(r: Vector[3], v: Vector[3], cd: Param, dt: Param, g: Param, terminated: Terminated,
r_next: Mutable[Vector[3]], v_next: Mutable[Vector[3]]):
a = vec(0.0, 0.0, -g) - cd * norm(v) * v
v_new = v + dt * a
v_next = v_new
r_next = r + dt * v_new
terminated = r_next[2] < 0.0
step = eagle_jax.function(flight, wrt=("r", "v", "cd"))
step
KernelFunction('flight', inputs=['r', 'v', 'cd', 'dt', 'g', 'terminated'], outputs=['r_next', 'v_next', 'terminated'], wrt=['r', 'v', 'cd'])
Scan the flight#
A finished ball’s outputs are zero (the kernel did not write them), so jnp.where keeps its last position. The scan
returns every position along the way and how many balls have landed after each step.
N, STEPS, DT, G, CD = 256, 360, 0.02, 9.81, 0.004
k = jax.random.split(jax.random.key(0), 3)
speed = 20 + 20 * jax.random.uniform(k[0], (N,))
elevation = jnp.radians(20) + jnp.radians(50) * jax.random.uniform(k[1], (N,))
heading = 2 * jnp.pi * jax.random.uniform(k[2], (N,))
v0 = jnp.stack([speed * jnp.cos(elevation) * jnp.cos(heading),
speed * jnp.cos(elevation) * jnp.sin(heading),
speed * jnp.sin(elevation)])
r0 = jnp.zeros((3, N)).at[2].set(1.0)
def body(carry, _, cd=CD):
r, v, landed = carry
r_next, v_next, landed_next = step(r, v, cd, DT, G, landed)
r_next, v_next = jnp.where(landed, r, r_next), jnp.where(landed, v, v_next)
return (r_next, v_next, landed_next), (r_next, landed_next.sum())
@jax.jit
def fly(r, v, landed):
return jax.lax.scan(body, (r, v, landed), None, length=STEPS)
(r, v, landed), (path, count) = fly(r0, v0, jnp.zeros(N, dtype=bool))
print("landed after", STEPS, "steps:", int(landed.sum()), "of", N)
print("landed after every 60th step:", count[59::60].tolist())
landed after 360 steps: 256 of 256
landed after every 60th step: [0, 27, 96, 190, 247, 256]
Left, the paths of 12 balls, each frozen where it landed; right, how many balls have landed as the flight goes on.
import matplotlib.pyplot as plt
fig, (ax, bx) = plt.subplots(1, 2, figsize=(9, 3.4))
for i, colour in zip(range(12), plt.cm.viridis([j / 11 for j in range(12)])):
ax.plot(path[:, 0, i], path[:, 2, i], color=colour, lw=1.2)
ax.set(xlabel="x (m)", ylabel="height (m)", title="12 of the 256 balls")
bx.plot(jnp.arange(1, STEPS + 1) * DT, count)
bx.set(xlabel="time (s)", ylabel="balls landed", title="the mask grows, never shrinks")
plt.tight_layout()
The same loop on every platform#
The kernel runs inline on the CPU and on XLA’s stream on a GPU. Put the inputs on each device JAX has and compare: the numbers match.
results = {}
for platform in ("cpu", "gpu"):
try:
device = jax.devices(platform)[0]
except RuntimeError:
continue # this JAX has no such device
inputs = [jax.device_put(x, device) for x in (r0, v0, jnp.zeros(N, dtype=bool))]
(r_p, _, landed_p), _ = fly(*inputs)
results[platform] = (int(landed_p.sum()), float(jnp.hypot(r_p[0], r_p[1]).mean()))
print(f"{platform}: {results[platform][0]} landed, mean range {results[platform][1]:.6f} m")
if len(results) == 2:
assert results["cpu"][0] == results["gpu"][0] and abs(results["cpu"][1] - results["gpu"][1]) < 1e-9
print("identical on both")
cpu: 256 landed, mean range 67.477918 m
gpu: 256 landed, mean range 67.477918 m
identical on both
Donate the mask#
Calling the stepping function from your own loop instead of a scan? Pass the old mask with donate_argnums and the
kernel updates it in place: the returned mask is the same buffer.
one_step = jax.jit(lambda r, v, landed: step(r, v, CD, DT, G, landed), donate_argnums=2)
mask = jnp.zeros(N, dtype=bool)
pointer = mask.unsafe_buffer_pointer()
_, _, mask_next = one_step(r0, v0, mask)
print("same buffer:", mask_next.unsafe_buffer_pointer() == pointer, "| the old handle is spent:", mask.is_deleted())
same buffer: True | the old handle is spent: True
A gradient through the whole loop#
How far do the balls travel, and how does that change with drag? grad runs backward through all 360 steps and
through the mask: a landed ball contributes nothing after it lands. The landing step is a switch, so this is the
slope with each ball’s landing step held fixed.
def mean_range(cd):
(r, _, _), _ = jax.lax.scan(lambda c, _: body(c, _, cd), (r0, v0, jnp.zeros(N, dtype=bool)), None, length=STEPS)
return jnp.hypot(r[0], r[1]).mean()
range_m, slope = jax.jit(jax.value_and_grad(mean_range))(CD)
print(f"mean range {float(range_m):.3f} m; d range / d cd = {float(slope):.1f} m")
mean range 67.478 m; d range / d cd = -2759.4 m
What just happened#
The kernel finished the balls; the mask rode through
scanas an ordinary carry.Eager,
jit,scanand a donated mask all take the same path, on the CPU and the GPU.gradwent backward through every step, stopping at each landing.
Next#
eagle inside your JAX program: what crosses, what is refused, the limits.
Gradients through a graph: the same idea for a multi-kernel step in eagle itself.