"""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)))