"""JAX Laplacian contraction against an independent Hessian trace."""
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np
from omnibias.jax.laplacian import deep_field_laplacian

weight = jnp.array([[0.2, -0.3], [0.5, 0.1]])
bias = jnp.array([0.1, -0.1])
readout = jnp.array([[0.7, -0.4]])
point = jnp.array([[0.3, -0.2]])
def lap(W):
    return deep_field_laplacian(point, [(W, bias, "tanh"), (readout, None, None)])
def field(x):
    return (readout @ jnp.tanh(weight @ x + bias)).squeeze()
reference = jnp.trace(jax.hessian(field)(point[0]))
np.testing.assert_allclose(lap(weight).squeeze(), reference, rtol=1e-10, atol=1e-12)
gradient = jax.grad(lambda W: jnp.mean(lap(W)**2))(weight)
assert gradient.shape == weight.shape and bool(jnp.isfinite(gradient).all())
print("JAX Laplacian and parameter gradient verified.")