From 0e0214860bd9d6368ee6305998b5fa1615881fea Mon Sep 17 00:00:00 2001 From: AHMETHAKANBEZIR1 Date: Wed, 21 Oct 2026 04:27:13 +0300 Subject: [PATCH] Select Subspace feature sequences with array indexing Co-authored-by: OpenAI Codex --- news/+subspace-feature-sequences.bugfix.rst | 2 ++ src/tinygp/transforms.py | 3 +- tests/test_subspace_sequences.py | 40 +++++++++++++++++++++ 3 files changed, 44 insertions(+), 1 deletion(-) create mode 100644 news/+subspace-feature-sequences.bugfix.rst create mode 100644 tests/test_subspace_sequences.py diff --git a/news/+subspace-feature-sequences.bugfix.rst b/news/+subspace-feature-sequences.bugfix.rst new file mode 100644 index 00000000..1f127347 --- /dev/null +++ b/news/+subspace-feature-sequences.bugfix.rst @@ -0,0 +1,2 @@ +Fix feature selection with tuple or list indices in the ``Subspace`` transform, +including reordered and negative feature indices. diff --git a/src/tinygp/transforms.py b/src/tinygp/transforms.py index 65fe1b38..df6aaf6e 100644 --- a/src/tinygp/transforms.py +++ b/src/tinygp/transforms.py @@ -159,4 +159,5 @@ class Subspace(Kernel): kernel: Kernel def evaluate(self, X1: JAXArray, X2: JAXArray) -> JAXArray: - return self.kernel.evaluate(X1[self.axis], X2[self.axis]) + axis = self.axis if isinstance(self.axis, int) else jnp.asarray(self.axis) + return self.kernel.evaluate(X1[axis], X2[axis]) diff --git a/tests/test_subspace_sequences.py b/tests/test_subspace_sequences.py new file mode 100644 index 00000000..abd6d336 --- /dev/null +++ b/tests/test_subspace_sequences.py @@ -0,0 +1,40 @@ +import jax +import jax.numpy as jnp +import pytest + +from tinygp import GaussianProcess, kernels, transforms +from tinygp.test_utils import assert_allclose + + +@pytest.mark.parametrize("axes", [(0, 2), (2, 0), (-1, 0), (0,)]) +@pytest.mark.parametrize("as_list", [False, True]) +def test_subspace_sequence_selects_features(axes, as_list): + x = jnp.arange(15.0).reshape(5, 3) / 10 + scale = jnp.arange(1, len(axes) + 1, dtype=float) + kernel = transforms.Subspace( + list(axes) if as_list else axes, + transforms.Linear(scale, kernels.ExpSquared()), + ) + + def reference(values): + selected = jnp.stack([values[:, axis] for axis in axes], axis=-1) * scale + differences = selected[:, None] - selected[None] + return jnp.exp(-0.5 * jnp.sum(differences**2, axis=-1)) + + assert_allclose(jax.jit(lambda values: kernel(values, values))(x), reference(x)) + assert_allclose(kernel(x), jnp.ones(len(x))) + actual_grad = jax.jit(jax.grad(lambda values: kernel(values, values).sum()))(x) + expected_grad = jax.jit(jax.grad(lambda values: reference(values).sum()))(x) + assert_allclose(actual_grad, expected_grad) + + +def test_subspace_sequence_in_gp_likelihood(): + x = jnp.arange(15.0).reshape(5, 3) / 10 + y = jnp.sin(x[:, 0]) + kernel = transforms.Subspace((0, 2), kernels.ExpSquared(0.7)) + selected = jnp.stack([x[:, 0], x[:, 2]], axis=-1) + covariance = kernels.ExpSquared(0.7)(selected, selected) + 0.1 * jnp.eye(len(x)) + expected = jax.scipy.stats.multivariate_normal.logpdf( + y, jnp.zeros_like(y), covariance + ) + assert_allclose(GaussianProcess(kernel, x, diag=0.1).log_probability(y), expected)