import functools
import numpy as np
import weanj # sets XLA_PYTHON_CLIENT_PREALLOCATE=false before JAX initialises
import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)
EXAMPLE, K_CHECK = "02-gauss-seidel", 20
d = weanj.load_inputs(EXAMPLE)
theta0, uobs0 = jnp.asarray(d["theta"]), jnp.asarray(d["uobs"])02. Gauss–Seidel — JAX
In [1]:
In [2]:
# Lexicographic Gauss–Seidel. Arrays are immutable and loop bodies must be pure, so
# the sweep is a scan over columns, each a scan over rows carrying the new u[i-1, j].
# Same update order as the Julia loop (i fastest, then j).
@functools.partial(jax.jit, static_argnums=2)
def gs_loss(theta, uobs, K):
n = theta.shape[0]
h2 = 1.0 / (n + 1) ** 2
def column(u, j):
left, old, right = u[:, j], u[:, j + 1], u[:, j + 2]
def point(below, i): # below = new u[i-1, j]
new = (h2 * theta[i, j] + below + old[i + 2] + left[i + 1] + right[i + 1]) / 4
return new, new
_, col = jax.lax.scan(point, old[0], jnp.arange(n))
return u.at[1:n + 1, j + 1].set(col), None
def sweep(u, _):
return jax.lax.scan(column, u, jnp.arange(n))[0], None
u, _ = jax.lax.scan(sweep, jnp.zeros((n + 2, n + 2), theta.dtype), None, length=K)
return jnp.sum((u[1:-1, 1:-1] - uobs) ** 2)In [3]:
gs_grad = jax.jit(jax.grad(gs_loss), static_argnums=2)
g_lex = gs_grad(theta0, uobs0, K_CHECK)
print(f"loss = {float(gs_loss(theta0, uobs0, K_CHECK)):.12f} ‖∇loss‖ = {float(jnp.linalg.norm(g_lex)):.6e}")loss = 0.105860373742 ‖∇loss‖ = 6.409859e-04In [4]:
# Red-black Gauss–Seidel: the reformulation JAX users actually write. Each colour is
# one vectorised stencil, selected with a mask.
@functools.partial(jax.jit, static_argnums=2)
def rb_loss(theta, uobs, K):
n = theta.shape[0]
h2 = 1.0 / (n + 1) ** 2
i, j = jnp.meshgrid(jnp.arange(2, n + 2), jnp.arange(2, n + 2), indexing="ij")
red = (i + j) % 2 == 0 # 1-based indices, as in the Julia kernel
def half(u, colour):
new = (h2 * theta + u[:-2, 1:-1] + u[2:, 1:-1] + u[1:-1, :-2] + u[1:-1, 2:]) / 4
return u.at[1:-1, 1:-1].set(jnp.where(colour, new, u[1:-1, 1:-1]))
def sweep(u, _):
return half(half(u, red), ~red), None
u, _ = jax.lax.scan(sweep, jnp.zeros((n + 2, n + 2), theta.dtype), None, length=K)
return jnp.sum((u[1:-1, 1:-1] - uobs) ** 2)In [5]:
rb_grad = jax.jit(jax.grad(rb_loss), static_argnums=2)
g_rb = rb_grad(theta0, uobs0, K_CHECK)
print(f"loss = {float(rb_loss(theta0, uobs0, K_CHECK)):.12f} ‖∇loss‖ = {float(jnp.linalg.norm(g_rb)):.6e}")loss = 0.105843754705 ‖∇loss‖ = 6.366981e-04In [6]:
e1 = weanj.fd_check(lambda t: gs_loss(jnp.asarray(t), uobs0, K_CHECK), np.asarray(theta0), np.asarray(g_lex), n=10, h=1e-2)
e2 = weanj.fd_check(lambda t: rb_loss(jnp.asarray(t), uobs0, K_CHECK), np.asarray(theta0), np.asarray(g_rb), n=10, h=1e-2)
weanj.save_gradient(EXAMPLE, "jax-lex", np.asarray(g_lex))
weanj.save_gradient(EXAMPLE, "jax-rb", np.asarray(g_rb))
print(f"finite differences: lexicographic {e1:.1e}, red-black {e2:.1e}")finite differences: lexicographic 1.1e-09, red-black 9.3e-10In [7]:
import time
ns, K, budget_s = [32, 64, 128, 256, 512, 1024], 100, 5.0
def inputs(n, seed=1):
rng = np.random.default_rng(seed)
return jnp.asarray(rng.standard_normal((n, n))), jnp.asarray(0.01 * rng.standard_normal((n, n)))
def temp_bytes(fn, *args):
"""XLA's scratch memory for the compiled gradient (excludes inputs/outputs)."""
return int(fn.lower(*args).compile().memory_analysis().temp_size_in_bytes)
def sweep(variant, loss, grad):
records = []
for n in ns:
theta, uobs = inputs(n)
t0 = time.perf_counter()
jax.block_until_ready(grad(theta, uobs, K))
compile_s = time.perf_counter() - t0
samples = 3 if n >= 512 else 10
r = dict(dtype="Float64", variant=variant, N=n * n, n=n, K=K,
forward_s=weanj.time_median(lambda: loss(theta, uobs, K), sync=jax.block_until_ready, samples=samples),
grad_s=weanj.time_median(lambda: grad(theta, uobs, K), sync=jax.block_until_ready, samples=samples),
compile_s=compile_s, peak_mem_bytes=temp_bytes(grad, theta, uobs, K))
records.append(r)
if r["grad_s"] > budget_s:
break
return records
for device in ("cpu", "gpu"):
with jax.default_device(jax.devices(device)[0]):
records = sweep("lexicographic", gs_loss, gs_grad) + sweep("red-black", rb_loss, rb_grad)
if device == "cpu": # gradient memory against the number of sweeps K, at n = 256
theta, uobs = inputs(256)
for Km in (10, 100, 1000):
records.append(dict(dtype="Float64", variant="lexicographic", N=256 * 256, n=256, K=Km,
forward_s=None, grad_s=None, compile_s=None,
peak_mem_bytes=temp_bytes(gs_grad, theta, uobs, Km)))
weanj.write_results(EXAMPLE, "jax", device, records)
for r in records:
if r["grad_s"] is not None:
print(f"{device} {r['variant']:<13} n={r['n']:<5d} forward {1e3*r['forward_s']:9.3f} ms "
f"gradient {1e3*r['grad_s']:9.3f} ms temp {r['peak_mem_bytes']/2**20:8.2f} MiB")cpu lexicographic n=32 forward 3.558 ms gradient 9.862 ms temp 0.04 MiB
cpu lexicographic n=64 forward 9.301 ms gradient 46.664 ms temp 0.14 MiB
cpu lexicographic n=128 forward 27.551 ms gradient 246.166 ms temp 0.52 MiB
cpu lexicographic n=256 forward 99.696 ms gradient 3718.064 ms temp 2.04 MiB
cpu lexicographic n=512 forward 517.822 ms gradient 19701.803 ms temp 8.07 MiB
cpu red-black n=32 forward 0.613 ms gradient 2.752 ms temp 0.07 MiB
cpu red-black n=64 forward 1.603 ms gradient 9.108 ms temp 0.27 MiB
cpu red-black n=128 forward 8.176 ms gradient 42.134 ms temp 1.06 MiB
cpu red-black n=256 forward 46.946 ms gradient 143.760 ms temp 4.19 MiB
cpu red-black n=512 forward 169.794 ms gradient 429.384 ms temp 16.63 MiB
cpu red-black n=1024 forward 724.900 ms gradient 1806.691 ms temp 66.25 MiB
gpu lexicographic n=32 forward 1068.302 ms gradient 1878.299 ms temp 0.04 MiB
gpu lexicographic n=64 forward 4231.997 ms gradient 7322.704 ms temp 0.11 MiB
gpu red-black n=32 forward 0.984 ms gradient 1.698 ms temp 0.03 MiB
gpu red-black n=64 forward 1.178 ms gradient 1.784 ms temp 0.11 MiB
gpu red-black n=128 forward 1.248 ms gradient 1.977 ms temp 0.42 MiB
gpu red-black n=256 forward 2.019 ms gradient 3.192 ms temp 3.13 MiB
gpu red-black n=512 forward 4.093 ms gradient 7.229 ms temp 12.52 MiB
gpu red-black n=1024 forward 9.076 ms gradient 19.968 ms temp 26.06 MiB