What a Weight Can Be, in JAX/Flax NNX

The explainer said a kernel is a price list for roughness: its eigenvalues set a cost for every frequency, and the RKHS is the weights you can afford. This is that, in JAX. The eigenvalues are the kernel’s spectral density, which a fast Fourier transform hands you for free; the RKHS norm of a weight is a single sum over them; and the same numbers settle where the Yat kernel actually lives. Every number is from a real run; the script is scripts/render_weight_be_gifs.py.

Prefer a notebook? Run the whole thing on Kaggle: the same code, block by block, with the math written out and every figure reproduced from a real run.

A kernel’s eigenvalues are its spectral density#

The whole price list is one FFT. On a periodic domain a translation-invariant kernel $k(x, y) = \kappa(x - y)$ is diagonal in the Fourier basis: its eigenfunctions are the modes $\cos(kx), \sin(kx)$, and its eigenvalues $\lambda_k$ are the Fourier transform of $\kappa$.

import jax, jax.numpy as jnp

N = 2048
grid = jnp.linspace(-jnp.pi, jnp.pi, N, endpoint=False)

@jax.jit
def spectrum(kappa):                                   # λₖ, the kernel's eigenvalues
    return jnp.fft.rfft(jnp.fft.ifftshift(kappa)).real

Feed it three kernels and the three personalities from the explainer fall straight out. The decays are known in closed form, and the FFT agrees with them down to floating-point precision:

sig, eps = 0.35, 0.12
lam_gauss = spectrum(jnp.exp(-grid**2 / (2*sig**2)))   # λₖ ∝ exp(−σ²k²/2)   super-exponential
lam_sobol = spectrum(jnp.exp(-jnp.abs(grid) / sig))    # λₖ ∝ (1 + σ²k²)⁻¹    polynomial ~ k⁻²
lam_imq   = spectrum(1.0 / (grid**2 + eps))            # λₖ ∝ exp(−√ε·k)      exponential

The Gaussian’s eigenvalues fall off a cliff, so its space holds only the smoothest, analytic weights. The Sobolev kernel’s fall like a gentle power, $\lambda_k \sim k^{-2}$, leaving room for rougher functions. The inverse-multiquadric, which is the Yat kernel’s distance gate, falls off exponentially: faster than the Sobolev kernel, slower than the Gaussian, sitting squarely between them.

The bill: which weights are affordable#

A weight written in the modes, $f = \sum_k c_k\,\phi_k$, has the RKHS norm

one charge per mode. A corner is a triangle wave, whose energy decays as $c_k^2 \sim k^{-4}$. Sum the charges and watch the bill.

import numpy as np
k = jnp.arange(1, 91)
ck2 = jnp.where(k % 2 == 1, (1.0 / k**2)**2, 0.0)      # a corner: cₖ² ~ 1/k⁴
ck2 = ck2 / ck2.sum()

def bill(lam_k):                                       # the running RKHS norm Σ cₖ²/λₖ
    return jnp.cumsum(ck2 / lam_k)

Under the Sobolev price list the charges shrink like $k^{-4} \cdot k^{2} = k^{-2}$, so the bill converges: the corner is in the space. Under the Gaussian the charges explode like $k^{-4}\,e^{\sigma^2 k^2 / 2}$, so the bill runs to infinity fast. And under the inverse-multiquadric the charges grow like $k^{-4}\,e^{\sqrt{\varepsilon}\,k}$: the bill still diverges, just slowly. A corner is not in the Yat kernel’s space either, only it takes many more modes to notice.

The same corner, fit in two homes#

You can feel the difference without any spectra. Fit a tent by kernel ridge, with a Gaussian, a Yat (inverse-multiquadric), and a Sobolev (Laplace) kernel, and lower the penalty. The Sobolev fit sharpens into the corner; both smooth kernels, the Gaussian and the Yat, round it off forever, because a corner is in neither of their spaces at any price. And notice that nothing here is trained: the weight is placed, the representer form written as a tiny Flax NNX module holding the centers and the solved coefficients.

from flax import nnx

class KernelWeight(nnx.Module):
    """A placed weight: centers X, coefficients α, read by the kernel. Nothing trained."""
    def __init__(self, X, alpha, kfn):
        self.X = nnx.Variable(X); self.alpha = nnx.Variable(alpha); self.kfn = kfn
    def __call__(self, xq):
        return self.kfn(xq[:, None], self.X.value[None, :]) @ self.alpha.value

def ridge(X, y, kfn, lam):                             # solve (G + λI) α = y, the representer weight
    G = kfn(X[:, None], X[None, :])
    alpha = jnp.linalg.solve(G + lam * jnp.eye(len(X)), y)
    return KernelWeight(X, alpha, kfn)

gauss = lambda a, b: jnp.exp(-((a - b)**2) / (2 * 0.14**2))
sobol = lambda a, b: jnp.exp(-jnp.abs(a - b) / 0.14)   # Laplace = H¹
yat   = lambda a, b: 1.0 / ((a - b)**2 + 0.03)         # inverse-multiquadric gate

On the sphere, the modes are harmonics#

Normalized data lives on a sphere, where the natural kernels depend only on the angle, $k(x, y) = \kappa(\langle x, y\rangle)$. Their eigenfunctions are the spherical harmonics, graded by degree, and the eigenvalue of a whole degree is its Legendre coefficient (the Funk-Hecke theorem):

def legendre(k, t):                                    # Pₖ by recurrence
    p0, p1 = jnp.ones_like(t), t
    for n in range(1, k):
        p0, p1 = p1, ((2*n+1) * t * p1 - n * p0) / (n + 1)
    return p1 if k > 0 else jnp.ones_like(t)

t = jnp.linspace(-1, 1, 400)
def degree_spectrum(kappa, K=10):                      # bₖ for a zonal kernel κ(t)
    return jnp.array([(2*k+1)/2 * jnp.trapezoid(kappa(t) * legendre(k, t), t) for k in range(K+1)])

Sharpen the cap, $\kappa(t) = e^{s(t - 1)}$, and its high degrees grow cheaper, so the same weight (a sum of caps at a few prototypes) is allowed to ripple faster across the globe.

Where the Yat kernel actually sits#

So which home does this series’ prototype move into? The explainer already placed it: not Sobolev, roomier than the Gaussian, universal for $\varepsilon > 0$, zonal on the sphere, its legibility bought by locality rather than roughness. What the code adds is the numbers behind that placement, the three decays side by side:

# the three decays, side by side: super-exponential, polynomial, exponential
for name, lam in [('Gaussian', lam_gauss), ('Sobolev', lam_sobol), ('Yat / IMQ', lam_imq)]:
    l = np.asarray(lam); l = l / l[1]
    print(f'{name:10s} λ₈={l[8]:.1e}  λ₁₆={l[16]:.1e}  λ₃₂={l[32]:.1e}')
# Gaussian   λ₈=2.1e-02  λ₁₆=1.6e-07  λ₃₂=3.7e-33     (a cliff: analytic only)
# Sobolev    λ₈=1.3e-01  λ₁₆=3.5e-02  λ₃₂=8.9e-03     (polynomial: corners allowed)
# Yat / IMQ  λ₈=8.7e-02  λ₁₆=5.4e-03  λ₃₂=2.5e-06     (exponential: smooth, universal, between)

Exponential decay, between the Gaussian’s cliff and the Sobolev power law: smooth, universal, roomier than the Gaussian, and not Sobolev, exactly where the explainer placed it.

Positive-definite kernels on the sphere are Schoenberg (1942); native spaces and the Sobolev correspondence are in Wendland (2004); the smallness of the Gaussian RKHS is in Steinwart and Christmann (2008); the polynomial-alignment and IMQ construction is Bouhsine (2026). The conceptual companion is What Can a Weight Be?.

@misc{bouhsine2026whatcanaweightbejaxflaxnnx,
  author       = {Bouhsine, Taha},
  title        = {What a Weight Can Be, in JAX/Flax NNX},
  year         = {2026},
  month        = {jun},
  howpublished = {\url{https://tahabouhsine.com/blog/what-can-a-weight-be-jax-flax-nnx/}},
  note         = {Blog post, Records of the !mmortal Data Scientist}
}