diff --git a/news/+traced-constant-sum.bugfix.rst b/news/+traced-constant-sum.bugfix.rst new file mode 100644 index 00000000..e2b1b2f9 --- /dev/null +++ b/news/+traced-constant-sum.bugfix.rst @@ -0,0 +1,3 @@ +Allow a traced scalar constant on the left of a kernel sum when using JIT, +automatic differentiation, or vectorization. Keep the builtin-zero shortcut +used by Python's ``sum``. diff --git a/src/tinygp/kernels/base.py b/src/tinygp/kernels/base.py index 8df5077b..0bfa6dfd 100644 --- a/src/tinygp/kernels/base.py +++ b/src/tinygp/kernels/base.py @@ -109,7 +109,7 @@ def __add__(self, other: Kernel | JAXArray) -> Kernel: def __radd__(self, other: Any) -> Kernel: # We'll hit this first branch when using the `sum` function - if other == 0: + if isinstance(other, (int, float)) and other == 0: return self if isinstance(other, Kernel): return Sum(other, self) diff --git a/tests/test_kernels/test_traced_constant_sum.py b/tests/test_kernels/test_traced_constant_sum.py new file mode 100644 index 00000000..6f45f14c --- /dev/null +++ b/tests/test_kernels/test_traced_constant_sum.py @@ -0,0 +1,49 @@ +import jax +import jax.numpy as jnp +import pytest + +from tinygp import GaussianProcess, kernels +from tinygp.test_utils import assert_allclose + + +@pytest.mark.parametrize("constant", [0.0, 0.5]) +def test_left_constant_sum_under_jax_transforms(constant): + x = jnp.linspace(-1, 1, 5) + kernel = kernels.ExpSquared(0.7) + + def evaluate(value): + return (value + kernel)(x, x) + + assert_allclose(jax.jit(evaluate)(constant), kernel(x, x) + constant) + derivative = jax.jit(jax.jacrev(evaluate))(constant) + assert_allclose(derivative, jnp.ones((len(x), len(x)))) + constants = jnp.array([0.0, 0.5, 1.0]) + expected = kernel(x, x)[None] + constants[:, None, None] + assert_allclose(jax.vmap(evaluate)(constants), expected) + + +def test_left_constant_sum_gp_likelihood_gradient(): + x = jnp.linspace(-1, 1, 5) + y = jnp.sin(x) + kernel = kernels.ExpSquared(0.7) + + def left(value): + return GaussianProcess(value + kernel, x, diag=0.1).log_probability(y) + + def dense_reference(value): + covariance = kernel(x, x) + value + 0.1 * jnp.eye(len(x)) + return jax.scipy.stats.multivariate_normal.logpdf( + y, jnp.zeros_like(y), covariance + ) + + actual = jax.jit(jax.value_and_grad(left))(0.5) + expected = jax.jit(jax.value_and_grad(dense_reference))(0.5) + assert_allclose(actual[0], expected[0]) + assert_allclose(actual[1], expected[1]) + + +def test_builtin_zero_sum_keeps_kernel_identity(): + kernel = kernels.ExpSquared(0.7) + assert 0 + kernel is kernel + assert 0.0 + kernel is kernel + assert sum([kernel]) is kernel