Batches and Jacobians you already use#
vmap, jacfwd, jacrev and hessian work on a hawk kernel as they do on any JAX function. Some of them cost
one launch, some one per batch entry. This page shows which, then draws a loss surface with them.
Time: ~5 min · Runs on: CPU, GPU if visible (picked up automatically) · You need: eagle, hawk and JAX
You will build: the launch count of each transformation, a Jacobian two ways, and a map of a two-Param loss with its gradient arrows and curvature.
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.
The kernel#
The same drag step as the training tutorial, with g a Param too, so there are two numbers to vary.
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, Vector
from hawk.math import norm, vec
from eagle.frameworks import jax as eagle_jax
@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
step = eagle_jax.function(drag_step, wrt=("v", "cd", "g"))
vmap over trajectories: one launch#
Eight batches of 256 balls. vmap folds the batch into the kernel’s sample axis, so the whole thing is one
launch. It is free for in_axes=-1; a leading batch axis costs one transpose.
batches = jax.random.normal(jax.random.key(0), (8, 3, 256))
over_axis0 = jax.vmap(lambda v: step(v, 0.05, 0.02, 9.81))
looped = jnp.stack([step(v, 0.05, 0.02, 9.81) for v in batches])
assert jnp.allclose(over_axis0(batches), looped)
print("vmap over 8 batches:", over_axis0(batches).shape, "| same as 8 separate calls")
vmap over 8 batches: (8, 3, 256) | same as 8 separate calls
vmap over Params: one launch too#
A batch of Param values is not a sample axis, so the kernel reads its own value from a table, one entry per
batch entry: one launch over all the samples of all the entries. Count the launches by reading the lowered
program: the call sites, and whether they sit in a loop.
v = batches[0]
cds = jnp.linspace(0.01, 0.2, 8)
cases = {
"one call": (lambda v: step(v, 0.05, 0.02, 9.81), (v,)),
"vmap over planes": (over_axis0, (batches,)),
"vmap over Params": (jax.vmap(lambda c: step(v, c, 0.02, 9.81)), (cds,)),
"grad": (jax.grad(lambda c: step(v, c, 0.02, 9.81).sum()), (0.05,)),
"jacrev over the planes": (jax.jacrev(lambda x: step(x, 0.05, 0.02, 9.81)), (v[:, :4],)),
"jacfwd over the planes": (jax.jacfwd(lambda x: step(x, 0.05, 0.02, 9.81)), (v[:, :4],)),
}
print(f"{'':<24}{'launch sites':<14}launches")
for name, (f, args) in cases.items():
text = jax.jit(f).lower(*args).as_text()
in_loop = "stablehlo.while" in text
print(f"{name:<24}{text.count('eagle_xla_launch'):<14}{'one per entry (loop)' if in_loop else 'one'}")
launch sites launches
one call 1 one
vmap over planes 1 one
vmap over Params 1 one
grad 1 one
jacrev over the planes 1 one
jacfwd over the planes 1 one
Jacobians, two ways#
How does one ball’s final velocity depend on its launch velocity? jacfwd and jacrev give the same 3 x 3
matrix, from the forward and the reverse kernel hawk derived.
def final_velocity(v_launch):
v = jnp.tile(v_launch[:, None], (1, 4)) # four copies: a kernel wants a batch
for _ in range(20):
v = step(v, 0.05, 0.02, 9.81)
return v[:, 0]
v_launch = jnp.array([6.0, 2.0, 8.0])
J_fwd, J_rev = jax.jacfwd(final_velocity)(v_launch), jax.jacrev(final_velocity)(v_launch)
assert jnp.allclose(J_fwd, J_rev)
print("d final velocity / d launch velocity:")
print(jnp.round(J_fwd, 4))
d final velocity / d launch velocity:
[[ 0.7894 -0.0201 -0.0591]
[-0.0201 0.843 -0.0197]
[-0.0603 -0.0201 0.7878]]
Forward mode through a stepping kernel#
A kernel that carries state and a stop mask differentiates in forward mode too. Here a ball falls with drag until
it lands; jvp and jacfwd through twenty steps agree with jacrev and with finite differences. A ball that
landed early took fewer steps, so its derivative with respect to cd is smaller.
from hawk import Scalar, Terminated
@hawk.kernel(steps=1)
def fall(h: Scalar, v: Scalar, cd: Param, dt: Param, terminated: Terminated,
h_next: Mutable[Scalar], v_next: Mutable[Scalar]):
v_next = v + dt * (cd * v * v - 9.81)
h_next = h + dt * v_next
terminated = h_next < 0.0
fall_step = eagle_jax.function(fall, wrt=("h", "v", "cd"))
def height_after(h0, cd, steps=20):
h, v, done = h0, jnp.zeros_like(h0), jnp.zeros(h0.shape, dtype=bool)
for _ in range(steps):
h_new, v_new, done_new = fall_step(h, v, cd, 0.02, done)
h, v, done = jnp.where(done, h, h_new), jnp.where(done, v, v_new), done_new
return h
h0 = jnp.array([0.3, 1.0, 5.0, 20.0]) # the first ball lands during the run
_, tangent = jax.jvp(lambda h: height_after(h, 0.05), (h0,), (jnp.ones(4),))
eps = 1e-6
differences = (height_after(h0 + eps, 0.05) - height_after(h0 - eps, 0.05)) / (2 * eps)
assert jnp.allclose(tangent, differences, atol=1e-7)
d_cd_fwd = jax.jacfwd(lambda c: height_after(h0, c))(0.05)
d_cd_rev = jax.jacrev(lambda c: height_after(h0, c))(0.05)
assert jnp.allclose(d_cd_fwd, d_cd_rev)
print("jvp matches finite differences; jacfwd matches jacrev")
print("d height / d cd per ball:", jnp.round(d_cd_fwd, 4))
jvp matches finite differences; jacfwd matches jacrev
d height / d cd per ball: [0.0261 0.1974 0.1974 0.1974]
A kernel with an active set#
A kernel built with Guard(active_set=True), the guard that lets eagle’s own launcher skip finished samples, also
works here. The bridge runs it on every sample, and grad and jvp read the same derived kernels as before.
from hawk.ext import Guard, Kind
from hawk.math import sqrt
@hawk.kernel(kind=Kind("tutorial_active", guard=Guard(active_set=True)))
def spring(w: Scalar, a: Param, terminated: Terminated, y: Mutable[Scalar]):
y = a * w * w + sqrt(w)
active = eagle_jax.function(spring)
w, none_finished = jnp.linspace(0.5, 2.0, 5), jnp.zeros(5, dtype=bool)
slope = jax.grad(lambda w: active(w, 0.7, none_finished).sum())(w)
assert jnp.allclose(slope, 1.4 * w + 0.5 / jnp.sqrt(w))
print("d y / d w:", jnp.round(slope, 4))
d y / d w: [1.4071 1.7595 2.1972 2.6672 3.1536]
A loss surface in two Params#
Fit the drag cd and gravity g from the final velocities of 256 balls (truth: cd = 0.07, g = 9.81). vmap
over a grid of Param pairs gives the loss and its gradient at every point; hessian gives the curvature at the minimum.
vs = 5.0 * jax.random.normal(jax.random.key(3), (3, 256))
def final(cd, g):
v = vs
for _ in range(20):
v = step(v, cd, 0.05, g)
return v
truth = final(0.07, 9.81)
loss = lambda cd, g: ((final(cd, g) - truth) ** 2).mean()
CD, GG = jnp.meshgrid(jnp.linspace(0.0, 0.14, 15), jnp.linspace(7.0, 12.0, 15))
value, (d_cd, d_g) = jax.vmap(jax.vmap(jax.value_and_grad(loss, argnums=(0, 1))))(CD, GG)
((h_cc, h_cg), (_, h_gg)) = jax.hessian(loss, argnums=(0, 1))(0.07, 9.81)
H = jnp.array([[h_cc, h_cg], [h_cg, h_gg]])
print("curvature at the minimum (eigenvalues):", jnp.round(jnp.linalg.eigvalsh(H), 2))
curvature at the minimum (eigenvalues): [2.80000e-01 1.05003e+03]
Contours of the loss with arrows pointing downhill (-grad, drawn at equal length). The valley is steep across cd and nearly flat along g, which is what the two curvatures above say: the data pins the drag down far better than gravity.
import matplotlib.pyplot as plt
from themed_figure import show
import numpy as np
fig, ax = plt.subplots(figsize=(6.2, 4.4))
ax.contourf(CD, GG, np.log10(value + 1e-12), levels=14, cmap="viridis", alpha=0.85)
span = np.array([0.14, 5.0]) # axis ranges, so arrows have equal length on the page
down = -np.stack([d_cd, d_g]) / span[:, None, None]
down = down / (np.linalg.norm(down, axis=0) + 1e-12) * 0.04
every = (slice(None, None, 2), slice(None, None, 2))
ax.quiver(CD[every], GG[every], down[0][every] * span[0], down[1][every] * span[1],
angles="xy", scale_units="xy", scale=1, color="white", width=0.004)
ax.plot(0.07, 9.81, "o", color="#d44")
ax.set(xlabel="drag coefficient cd", ylabel="gravity g", title="log10 loss and the downhill direction (red: the truth)")
plt.tight_layout()
show()
What just happened#
Planes and Params both batch as one launch: the table above reads it off the program.
jacfwdandjacrevuse the forward and reverse kernels hawk derived;hessianuses both.A stepping kernel with a stop mask, and one built with
Guard(active_set=True), differentiate in forward mode too.You never wrote a derivative.
Next#
A stepping loop in
lax.scan: balls that land and stop.eagle inside your JAX program: what crosses, what is refused, the limits.