implement custom_root as a HiJAX primitive with deferred derivatives
The deferred derivatives save trace time. We paid +50 LoC to introduce a
_RootLinearMap hijax primitive (and its rules) so as to avoid tracing
the derivative of the user-supplied `f` when only doing primal
evaluation. The new primitive represents L(p, x, v) = D_x f(p, x) v, so
it is linear in v. We preserve symbolic zeros where possible.
We paid another +10 LoC in CustomRoot.batch_dim_rule to implement logic
that custom_jvp used to provide for us:
```python
@jax.custom_jvp
def h(x):
return 0.
@h.defjvp
def h_jvp(primals, tangents):
x, = primals
dx, = tangents
return 0., x * dx
jax.vmap(lambda x: jax.jvp(h, (x,), (1.,)))(jnp.arange(3.))
# (array([0., 0., 0.]), array([0., 1., 2.]))
```
Notice that the primal output of `h_jvp` is unbatched while the tangent
output is batched. The custom_jvp machinery would merge primal and
tangent output batch dimensions and broadcast the primal as necessary,
so the author of the above code need not worry about it. With HiPrim
we're relying on VmapOf, which requires a batch_dim_rule.
Synthetic CPU benchmark:
```python
import gc
import statistics
import time
import jax
import jax.numpy as jnp
def root(a):
def f(x):
for _ in range(128):
x = x + 0.01 * jnp.sin(x)
return x - a
def solve(f, x):
def step(state):
i, x = state
y, dy = jax.jvp(f, (x,), (jnp.ones_like(x),))
return i + 1, x - y / dy
return jax.lax.while_loop(lambda s: s[0] < 8, step, (0, x))[1]
return jax.lax.custom_root(
f, jnp.ones_like(a), solve,
lambda g, b: b / g(jnp.ones_like(b)))
spec = jax.ShapeDtypeStruct((32,), jnp.float32)
a = jnp.full(spec.shape, 0.5, spec.dtype)
grad = jax.grad(lambda a: root(a).sum())
for name, fn, tracing in [
('primal trace', jax.jit(root), True),
('grad trace', jax.jit(grad), True),
('eager grad', grad, False),
]:
jax.clear_caches()
samples = []
for i in range(14):
if tracing:
jax.clear_caches()
gc.collect()
gc.disable()
try:
start = time.perf_counter()
if tracing:
result = fn.trace(spec).lojax
else:
result = jax.block_until_ready(fn(a))
elapsed = time.perf_counter() - start
finally:
gc.enable()
if i >= 2:
samples.append(1000 * elapsed)
print(f'{name}: {statistics.median(samples):.2f} ms')
```
| Measurement | Before | After |
| --- | ---: | ---: |
| Primal tracing, including lojax conversion | 44.7 ms | 14.8 ms |
| Grad tracing, including lojax conversion | 58.1 ms | 20.5 ms |
| Eager grad | 188 ms | 155 ms |
Tiny workloads got small regressions, like _0.5ms for grad tracing and
+0.8ms for eager grad.
The generated HLO is not universally identical because of symbolic zero
handling changes, and some regrouping of additions. My
astronomy-named friend did very extensive experimentation to validate
it.
jax
98f33a51 - implement custom_root as a HiJAX primitive with deferred derivatives
