Your first closed-form Laplacian

In this tutorial you will build a one-layer scalar field on ℝ^D, compute its closed-form Laplacian, and verify it agrees with autodiff to machine precision. About 5 minutes.

Prerequisites

pip install omnibias-jax

Step 1 — define a one-layer field

A one-layer scalar field is f(x) = Σ_h c_h·σ(W_h·x + β_h) + b.

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
D, H = 8, 32
kW, kb, kc, kx = jax.random.split(key, 4)

W    = jax.random.normal(kW, (H, D)) / jnp.sqrt(D)
beta = jax.random.normal(kb, (H,))
c    = jax.random.normal(kc, (H,))
b    = 0.0
x    = jax.random.normal(kx, (D,))

Step 2 — value, gradient, Laplacian in one call

from omnibias.jax import neural_field_value_grad_laplacian

val, grad, lap = neural_field_value_grad_laplacian(x, W, beta, c, b, "tanh")
print("f(x)      =", val)
print("∇f(x)     =", grad.shape)   # (D,)
print("∇²f(x)    =", lap)          # scalar

The Laplacian here is a contraction Σ_h c_h·σ''(z_h)·‖W_h‖². The dimension- dependent ‖W_h‖² is computed once and reused — the overhead is O(1) in D.

Step 3 — verify against autodiff

The whole point is that the closed form matches autodiff exactly. The Laplacian is the trace of the Hessian:

def f(x):
    z = W @ x + beta
    return jnp.sum(c * jnp.tanh(z)) + b

H_autodiff = jax.hessian(f)(x)
lap_autodiff = jnp.trace(H_autodiff)

print("closed form :", lap)
print("autodiff    :", lap_autodiff)
print("abs diff    :", abs(lap - lap_autodiff))   # ~1e-15 or better

What you just proved

The closed-form Laplacian is bit-for-bit the same answer as autodiff — it is not an approximation. It is simply computed without building (or differentiating through) a Hessian graph, which is why it is faster and uses far less memory.

Step 4 — go high-order

For a 4th-order (biharmonic) operator, use the poly-Laplacian. Note that it stays flat in k — no nested-autodiff blow-up:

from omnibias.jax import neural_field_polylaplacian

d2 = neural_field_polylaplacian(x, W, beta, c, b, "tanh", k=1)  # Laplacian
d4 = neural_field_polylaplacian(x, W, beta, c, b, "tanh", k=2)  # biharmonic Δ²

Where to go next

• PINN heat equation — use this Laplacian in a real PDE loss.
• Choosing an activation — when to use gaussian vs tanh.
• Performance guide — measure the O(1)-in-D win on your own sizes.