From 0a8c875c9fc6dcf0ea0c81f4fdf1e31e24f32b4a Mon Sep 17 00:00:00 2001 From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com> Date: Thu, 8 Oct 2026 02:29:34 +0530 Subject: [PATCH] Fix NaN gradients of the CARMA kernel --- news/228.bugfix | 1 + src/tinygp/kernels/quasisep.py | 21 +++++++++++++++++---- tests/test_kernels/test_quasisep.py | 29 +++++++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 4 deletions(-) create mode 100644 news/228.bugfix diff --git a/news/228.bugfix b/news/228.bugfix new file mode 100644 index 00000000..4516244d --- /dev/null +++ b/news/228.bugfix @@ -0,0 +1 @@ +Fixed NaN gradients of the `CARMA` kernel with respect to its parameters, caused by square roots of terms that are exactly zero. diff --git a/src/tinygp/kernels/quasisep.py b/src/tinygp/kernels/quasisep.py index f28ca81b..3082bf5c 100644 --- a/src/tinygp/kernels/quasisep.py +++ b/src/tinygp/kernels/quasisep.py @@ -773,7 +773,9 @@ def __init__(self, alpha: Any, beta: Any): _complex_select = _complex_mask * complex_idx % 2 # Construct the obsservation model => real + complex - om_real = jnp.sqrt(jnp.abs(acf.real)) + # (_safe_sqrt keeps the gradient finite for terms that are exactly zero, + # including the ones discarded by jnp.where below) + om_real = _safe_sqrt(jnp.where(_real_mask, jnp.abs(acf.real), 0.0)) a, b, c, d = ( 2 * acf.real, @@ -785,10 +787,11 @@ def __init__(self, alpha: Any, beta: Any): d2 = jnp.square(d) s2 = c2 + d2 denom = jnp.where(_real_mask, 1.0, 2 * c * s2) - h2_2 = d2 * (a * c - b * d) / denom - h2 = jnp.sqrt(h2_2) + h2_2 = jnp.where(_real_mask, 0.0, d2 * (a * c - b * d) / denom) + h2 = _safe_sqrt(h2_2) denom = jnp.where(_real_mask, 1.0, d) - h1 = (c * h2 - jnp.sqrt(a * d2 - s2 * h2_2)) / denom + rad = jnp.where(_real_mask, 0.0, a * d2 - s2 * h2_2) + h1 = (c * h2 - _safe_sqrt(rad)) / denom om_complex = jnp.array([h1, h2]) # for complex roots, every conjugate pair match one full celerite term, @@ -900,6 +903,16 @@ def transition_matrix(self, X1: JAXArray, X2: JAXArray) -> JAXArray: return tm_real + tm_complex_diag + -tm_complex_u.T + tm_complex_u +def _safe_sqrt(x: JAXArray) -> JAXArray: + """Square root that is 0, with a zero derivative, for non-positive input + + The derivative of ``jnp.sqrt`` at 0 is infinite, which turns gradients into + NaN when a term is exactly zero or is discarded by ``jnp.where``. + """ + positive = x > 0 + return jnp.where(positive, jnp.sqrt(jnp.where(positive, x, 1.0)), 0.0) + + @jax.jit def carma_roots(poly_coeffs: JAXArray) -> JAXArray: roots = jnp.roots(poly_coeffs[::-1], strip_zeros=False) diff --git a/tests/test_kernels/test_quasisep.py b/tests/test_kernels/test_quasisep.py index c585f989..993e52a3 100644 --- a/tests/test_kernels/test_quasisep.py +++ b/tests/test_kernels/test_quasisep.py @@ -2,6 +2,7 @@ import jax.numpy as jnp import jax.scipy as jsp import pytest +from jax.test_util import check_grads from numpy import random as np_random from tinygp import GaussianProcess @@ -138,6 +139,34 @@ def loss(params): loss(params) +@pytest.mark.parametrize( + "alpha, beta", + [ + # real root + ([0.01], [0.1]), + # complex roots, the second entry of the observation model is exactly 0 + ([0.5, 0.1], [0.1]), + # complex roots + ([0.1, 1.0], [0.1, 0.05]), + # real and complex roots + ([0.5, 1.5, 1.2], [0.3, 0.2]), + ], +) +def test_carma_grad(data, alpha, beta): + x, y, _ = data + + def loss(alpha, beta): + kernel = quasisep.CARMA(alpha=alpha, beta=beta) + return -GaussianProcess(kernel, x, diag=0.01).log_probability(y) + + with jax.enable_x64(True): # type: ignore + alpha = jnp.array(alpha, dtype=jnp.float64) + beta = jnp.array(beta, dtype=jnp.float64) + grads = jax.grad(loss, argnums=(0, 1))(alpha, beta) + assert all(jnp.all(jnp.isfinite(g)) for g in grads) + check_grads(loss, (alpha, beta), order=1, modes=["rev"], rtol=1e-4) + + def test_carma_quads(): alpha = jnp.array([1.4, 2.3, 1.5]) beta = jnp.array([0.1, 0.5])