Graph kernel consistency with other kernels (#560)

* toggle dark theme

* graph variational families

* variational family not working and example

* graph svgp fix

* more updates

* dtype problem

* kernel checks in variational family

* removed graph gaussian

* formatting issues

* ran uv run poe format to fix

* passing tests and formatting weirdness

* subclassing and tests

* segregated graph tests that still fail

* graph tests, changes to parent variational class signatures and graph kernel to allow SVGP shapes and dtypes

* docstring for new class

* remove ensure_2d function

* tolerant dimensions

* graph kernel consistent with other kernels

* meaningful name and circular import issue

* version bump and typechecking magic

Author: syedzayyan

gpjax/__init__.py

@@ -40,7 +40,7 @@
 __description__ = "Gaussian processes in JAX and Flax"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.13.0"
+__version__ = "0.13.1"
 
 __all__ = [
     "gps",

gpjax/kernels/computations/eigen.py

@@ -15,7 +15,6 @@
 
 
 import beartype.typing as tp
-import jax.numpy as jnp
 from jaxtyping import (
     Float,
     Num,
@@ -39,17 +38,4 @@ class EigenKernelComputation(AbstractKernelComputation):
     def _cross_covariance(
         self, kernel: Kernel, x: Num[Array, "N D"], y: Num[Array, "M D"]
     ) -> Float[Array, "N M"]:
-        # Transform the eigenvalues of the graph Laplacian according to the
-        # RBF kernel's SPDE form.
-        S = jnp.power(
-            kernel.eigenvalues
-            + 2
-            * kernel.smoothness.value
-            / kernel.lengthscale.value
-            / kernel.lengthscale.value,
-            -kernel.smoothness.value,
-        )
-        S = jnp.multiply(S, kernel.num_vertex / jnp.sum(S))
-        # Scale the transform eigenvalues by the kernel variance
-        S = jnp.multiply(S, kernel.variance.value)
-        return kernel(x, y, S=S)
+        return kernel(x, y)

gpjax/kernels/non_euclidean/graph.py

@@ -25,7 +25,10 @@
     AbstractKernelComputation,
     EigenKernelComputation,
 )
-from gpjax.kernels.non_euclidean.utils import jax_gather_nd
+from gpjax.kernels.non_euclidean.utils import (
+    calculate_heat_semigroup,
+    jax_gather_nd,
+)
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.parameters import (
     Parameter,
@@ -98,14 +101,12 @@ def __init__(
 
         super().__init__(active_dims, lengthscale, variance, n_dims, compute_engine)
 
-    def __call__(  # TODO not consistent with general kernel interface
+    def __call__(
         self,
         x: Int[Array, "N 1"],
         y: Int[Array, "M 1"],
-        *,
-        S,
-        **kwargs,
     ):
+        S = calculate_heat_semigroup(self)
         Kxx = (jax_gather_nd(self.eigenvectors, x) * S.squeeze()) @ jnp.transpose(
             jax_gather_nd(self.eigenvectors, y)
         )  # shape (n,n)

gpjax/kernels/non_euclidean/utils.py

@@ -13,13 +13,20 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
+
+import beartype.typing as tp
+import jax.numpy as jnp
 from jaxtyping import (
     Float,
     Int,
 )
 
 from gpjax.typing import Array
 
+if tp.TYPE_CHECKING:
+    from gpjax.kernels.non_euclidean.graph import GraphKernel
+
 
 def jax_gather_nd(
     params: Float[Array, " N *rest"], indices: Int[Array, " M 1"]
@@ -41,3 +48,26 @@ def jax_gather_nd(
     """
     tuple_indices = tuple(indices[..., i] for i in range(indices.shape[-1]))
     return params[tuple_indices]
+
+
+def calculate_heat_semigroup(kernel: GraphKernel) -> Float[Array, "N M"]:
+    r"""Returns the rescaled heat semigroup, S
+
+    Args:
+        kernel: instance of the graph kernel
+
+    Returns:
+        S
+    """
+    S = jnp.power(
+        kernel.eigenvalues
+        + 2
+        * kernel.smoothness.value
+        / kernel.lengthscale.value
+        / kernel.lengthscale.value,
+        -kernel.smoothness.value,
+    )
+    S = jnp.multiply(S, kernel.num_vertex / jnp.sum(S))
+    # Scale the transform eigenvalues by the kernel variance
+    S = jnp.multiply(S, kernel.variance.value)
+    return S