JAX: forward spatial derivatives, ordinary parameter gradients
Install omnibias-jax==0.5.0rc1. Enable x64 before constructing arrays for a float64 numerical comparison.
high_order_jax.py
"""JAX directional jets and ordinary parameter autodiff."""
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from omnibias.jax.jet import jet_to_tower, mlp_jet
def residual_loss(weight):
layers = [(weight, jnp.array([0.1, 0.2]), "tanh"),
(jnp.array([[0.6, -0.4]]), None, None)]
tower = jet_to_tower(mlp_jet(jnp.array([0.3]), jnp.ones(1), layers, order=4))
assert tower.shape == (5, 1)
return jnp.mean((tower[4] - 24.0) ** 2)
weight = jnp.array([[0.5], [-0.3]])
loss, gradient = jax.value_and_grad(residual_loss)(weight)
assert gradient.shape == weight.shape
assert bool(jnp.all(jnp.isfinite(gradient)))
print("gradient:", gradient.shape, "finite:", bool(jnp.isfinite(loss)))
mlp_jet computes derivatives along the supplied direction. jax.value_and_grad acts on the residual loss with respect to the weight matrix. The gradient has the same shape as the weights.
See JAX jets and the high-order residual walkthrough.