Does JaxGP support changepoint kernels?

[jakeyeung]

Hi GPJax enthusiasts,

Does GPJax support custom kernels where the distance function depends on different locations of the data?

For example, GPFlow supports change-points : https://gpflow.github.io/GPflow/develop/notebooks/getting_started/kernels.html

This allows custom kernel distances depending on where you are in the data.

When I tried to create custom kernels with GPJax, the __call__() function seems to take only two pairs of points as input, so the function wouldn't know when to switch to a different distance function.

If the input were two vectors, then you could in principle calculate different distance functions for different parts of the data. But is this easily doable with GPJax?

Here is an illustrative example from GPFlow documentation:

EDIT:

I tried defining two different pairwise distance that switches depending on whether the pair of points are left, right, or straddle a changepoint x*. I use "if" conditionals to define different regions with distinct covariance functions, e.g.:

from jax.config import config

config.update("jax_enable_x64", True)

from dataclasses import dataclass
from typing import Dict

from jax import jit
import jax.numpy as jnp
import jax.random as jr
import jax.scipy as jscipy

# import types
from beartype.typing import Union
from jaxtyping import Float
from gpjax.typing import (
    Array,
    ScalarFloat,
)
from simple_pytree import static_field
import tensorflow_probability.substrates.jax as tfp
import gpjax as gpx
from gpjax.base.param import param_field


@dataclass
class SharedKernel(gpx.kernels.AbstractKernel):
    r"""Kernel that is shared between first half and second half of data, but no correlation between first and second half of data."""
    lengthscale: Union[ScalarFloat, Float[Array, " D"]] = param_field(
        jnp.array(1.0), bijector=tfb.Softplus()
    )
    variance: ScalarFloat = param_field(
            jnp.array(1.0), bijector=tfb.Softplus()
    )
    tmax: float = 1.0
    name: str = "Shared"
    def __call__(self, x1: Float[Array, " D"], x2: Float[Array, " D"]) -> ScalarFloat:
        l = self.lengthscale
        v = self.variance
        tmax = self.tmax

        if x1 <= xstar and x2 > xstar:
            K = 0 # no correlation between first half of data and second
        else:
            K = v * jnp.exp(-((x1 - x2) / l)**2)
        return K.squeeze()

But I am getting TracerBoolConversionError errors, would appreciate if there's any insight on these custom kernels.

Jake

[jakeyeung]

I fixed the bug.

In the end I replaced the series of if loops with a jax.lax.switch() and it worked.

Jake

[thomaspinder]

Nice work @jakeyeung .

Would you like to open a PR to the codebase with your above code? I'd be happy to support a changepoint kernel.