gpjax 0.13.0__tar.gz → 0.13.1__tar.gz

{gpjax-0.13.0 → gpjax-0.13.1}/.github/workflows/build_docs.yml

@@ -40,7 +40,7 @@ jobs:
 
       # Install katex for math support
       - name: Install NPM
-        uses: actions/setup-node@v5
+        uses: actions/setup-node@v6
         with:
           node-version: 16
       - name: Install KaTeX
@@ -49,7 +49,7 @@ jobs:
 
       # Install uv
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           version: "latest"
 

{gpjax-0.13.0 → gpjax-0.13.1}/.github/workflows/integration.yml

@@ -27,7 +27,7 @@ jobs:
 
       # Install uv
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           version: "latest"
 

{gpjax-0.13.0 → gpjax-0.13.1}/.github/workflows/pr_greeting.yml

@@ -13,8 +13,8 @@ jobs:
     steps:
       - uses: actions/first-interaction@v3.1.0
         with:
-          repo-token: ${{ secrets.GITHUB_TOKEN }}
-          pr-message: >+
+          repo_token: ${{ secrets.GITHUB_TOKEN }}
+          pr_message: >+
             Thank you for opening your first PR into GPJax!
 
 

{gpjax-0.13.0 → gpjax-0.13.1}/.github/workflows/release.yml

@@ -80,7 +80,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
 
       - name: Install dependencies
         run: |
@@ -116,7 +116,7 @@ jobs:
           python-version: '3.11'
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
 
       - name: Install dependencies
         run: |
@@ -132,7 +132,7 @@ jobs:
           uv run bandit -r gpjax/ -f json -o bandit-report.json || echo "Bandit scan completed with warnings"
 
       - name: Upload security reports
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v5
         with:
           name: security-reports
           path: |
@@ -154,7 +154,7 @@ jobs:
           python-version: '3.11'
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
 
       - name: Build package
         run: |
@@ -166,7 +166,7 @@ jobs:
           uv run twine check dist/*
 
       - name: Upload build artifacts
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v5
         with:
           name: dist-packages
           path: dist/
@@ -264,7 +264,7 @@ jobs:
         uses: actions/checkout@v5
 
       - name: Download build artifacts
-        uses: actions/download-artifact@v5
+        uses: actions/download-artifact@v6
         with:
           name: dist-packages
           path: dist/
@@ -294,7 +294,7 @@ jobs:
 
     steps:
       - name: Download build artifacts
-        uses: actions/download-artifact@v5
+        uses: actions/download-artifact@v6
         with:
           name: dist-packages
           path: dist/

{gpjax-0.13.0 → gpjax-0.13.1}/.github/workflows/security-analysis.yml

@@ -28,7 +28,7 @@ jobs:
         python-version: '3.11'
 
     - name: Install uv
-      uses: astral-sh/setup-uv@v6
+      uses: astral-sh/setup-uv@v7
       with:
         version: "latest"
 
@@ -47,7 +47,7 @@ jobs:
         uv run bandit -r gpjax/ -f json -o bandit-report.json || true
 
     - name: Upload dependency scan results
-      uses: actions/upload-artifact@v4
+      uses: actions/upload-artifact@v5
       if: always()
       with:
         name: security-scan-results

{gpjax-0.13.0 → gpjax-0.13.1}/.github/workflows/test_docs.yml

@@ -32,7 +32,7 @@ jobs:
 
       # Install uv
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           version: "latest"
 

{gpjax-0.13.0 → gpjax-0.13.1}/.github/workflows/tests.yml

@@ -26,7 +26,7 @@ jobs:
 
       # Install uv
       - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
         with:
           version: "latest"
 

{gpjax-0.13.0 → gpjax-0.13.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.13.0
+Version: 0.13.1
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues

{gpjax-0.13.0 → gpjax-0.13.1}/gpjax/__init__.py

@@ -40,7 +40,7 @@ __license__ = "MIT"
 __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-0.13.0 → gpjax-0.13.1}/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-0.13.0 → gpjax-0.13.1}/gpjax/kernels/non_euclidean/graph.py

@@ -25,7 +25,10 @@ from gpjax.kernels.computations import (
     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 @@ class GraphKernel(StationaryKernel):
 
         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, "N 1"],
-        *,
-        S,
-        **kwargs,
+        y: Int[Array, "M 1"],
     ):
+        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-0.13.0 → gpjax-0.13.1}/gpjax/kernels/non_euclidean/utils.py

@@ -13,6 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
+
+import beartype.typing as tp
+import jax.numpy as jnp
 from jaxtyping import (
     Float,
     Int,
@@ -20,6 +24,9 @@ from jaxtyping import (
 
 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

{gpjax-0.13.0 → gpjax-0.13.1}/gpjax/variational_families.py

@@ -19,7 +19,10 @@ import beartype.typing as tp
 from flax import nnx
 import jax.numpy as jnp
 import jax.scipy as jsp
-from jaxtyping import Float
+from jaxtyping import (
+    Float,
+    Int,
+)
 
 from gpjax.dataset import Dataset
 from gpjax.distributions import GaussianDistribution
@@ -108,6 +111,7 @@ class AbstractVariationalGaussian(AbstractVariationalFamily[L]):
         self,
         posterior: AbstractPosterior[P, L],
         inducing_inputs: tp.Union[
+            Int[Array, "N D"],
             Float[Array, "N D"],
             Real,
         ],
@@ -140,7 +144,7 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
     def __init__(
         self,
         posterior: AbstractPosterior[P, L],
-        inducing_inputs: Float[Array, "N D"],
+        inducing_inputs: tp.Union[Int[Array, "N D"], Float[Array, "N D"]],
         variational_mean: tp.Union[Float[Array, "N 1"], None] = None,
         variational_root_covariance: tp.Union[Float[Array, "N N"], None] = None,
         jitter: ScalarFloat = 1e-6,
@@ -156,6 +160,12 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         self.variational_mean = Real(variational_mean)
         self.variational_root_covariance = LowerTriangular(variational_root_covariance)
 
+    def _fmt_Kzt_Ktt(self, Kzt, Ktt):
+        return Kzt, Ktt
+
+    def _fmt_inducing_inputs(self):
+        return self.inducing_inputs.value
+
     def prior_kl(self) -> ScalarFloat:
         r"""Compute the prior KL divergence.
 
@@ -178,7 +188,7 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         # Unpack variational parameters
         variational_mean = self.variational_mean.value
         variational_sqrt = self.variational_root_covariance.value
-        inducing_inputs = self.inducing_inputs.value
+        inducing_inputs = self._fmt_inducing_inputs()
 
         # Unpack mean function and kernel
         mean_function = self.posterior.prior.mean_function
@@ -202,7 +212,9 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
 
         return q_inducing.kl_divergence(p_inducing)
 
-    def predict(self, test_inputs: Float[Array, "N D"]) -> GaussianDistribution:
+    def predict(
+        self, test_inputs: tp.Union[Int[Array, "N D"], Float[Array, "N D"]]
+    ) -> GaussianDistribution:
         r"""Compute the predictive distribution of the GP at the test inputs t.
 
         This is the integral $q(f(t)) = \int p(f(t)\mid u) q(u) \mathrm{d}u$, which
@@ -222,7 +234,7 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         # Unpack variational parameters
         variational_mean = self.variational_mean.value
         variational_sqrt = self.variational_root_covariance.value
-        inducing_inputs = self.inducing_inputs.value
+        inducing_inputs = self._fmt_inducing_inputs()
 
         # Unpack mean function and kernel
         mean_function = self.posterior.prior.mean_function
@@ -241,6 +253,8 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         Kzt = kernel.cross_covariance(inducing_inputs, test_points)
         test_mean = mean_function(test_points)
 
+        Kzt, Ktt = self._fmt_Kzt_Ktt(Kzt, Ktt)
+
         # Lz⁻¹ Kzt
         Lz_inv_Kzt = solve(Lz, Kzt)
 
@@ -259,8 +273,10 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
             - jnp.matmul(Lz_inv_Kzt.T, Lz_inv_Kzt)
             + jnp.matmul(Ktz_Kzz_inv_sqrt, Ktz_Kzz_inv_sqrt.T)
         )
+
         if hasattr(covariance, "to_dense"):
             covariance = covariance.to_dense()
+
         covariance = add_jitter(covariance, self.jitter)
         covariance = Dense(covariance)
 
@@ -269,6 +285,53 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         )
 
 
+class GraphVariationalGaussian(VariationalGaussian[L]):
+    r"""A variational Gaussian defined over graph-structured inducing inputs.
+
+    This subclass adapts the :class:`VariationalGaussian` family to the
+    case where the inducing inputs are discrete graph node indices rather
+    than continuous spatial coordinates.
+
+    The main differences are:
+      * Inducing inputs are integer node IDs.
+      * Kernel matrices are ensured to be dense and 2D.
+    """
+
+    def __init__(
+        self,
+        posterior: AbstractPosterior[P, L],
+        inducing_inputs: Int[Array, "N D"],
+        variational_mean: tp.Union[Float[Array, "N 1"], None] = None,
+        variational_root_covariance: tp.Union[Float[Array, "N N"], None] = None,
+        jitter: ScalarFloat = 1e-6,
+    ):
+        super().__init__(
+            posterior,
+            inducing_inputs,
+            variational_mean,
+            variational_root_covariance,
+            jitter,
+        )
+        self.inducing_inputs = self.inducing_inputs.value.astype(jnp.int64)
+
+    def _fmt_Kzt_Ktt(self, Kzt, Ktt):
+        Ktt = Ktt.to_dense() if hasattr(Ktt, "to_dense") else Ktt
+        Kzt = Kzt.to_dense() if hasattr(Kzt, "to_dense") else Kzt
+        Ktt = jnp.atleast_2d(Ktt)
+        Kzt = (
+            jnp.transpose(jnp.atleast_2d(Kzt)) if Kzt.ndim < 2 else jnp.atleast_2d(Kzt)
+        )
+        return Kzt, Ktt
+
+    def _fmt_inducing_inputs(self):
+        return self.inducing_inputs
+
+    @property
+    def num_inducing(self) -> int:
+        """The number of inducing inputs."""
+        return self.inducing_inputs.shape[0]
+
+
 class WhitenedVariationalGaussian(VariationalGaussian[L]):
     r"""The whitened variational Gaussian family of probability distributions.
 
@@ -811,6 +874,7 @@ __all__ = [
     "AbstractVariationalFamily",
     "AbstractVariationalGaussian",
     "VariationalGaussian",
+    "GraphVariationalGaussian",
     "WhitenedVariationalGaussian",
     "NaturalVariationalGaussian",
     "ExpectationVariationalGaussian",

{gpjax-0.13.0 → gpjax-0.13.1}/tests/test_variational_families.py

@@ -25,6 +25,8 @@ from jaxtyping import (
     Array,
     Float,
 )
+import networkx as nx
+import numpy as np
 import numpyro.distributions as npd
 from numpyro.distributions import Distribution as NumpyroDistribution
 import pytest
@@ -35,6 +37,7 @@ from gpjax.variational_families import (
     AbstractVariationalFamily,
     CollapsedVariationalGaussian,
     ExpectationVariationalGaussian,
+    GraphVariationalGaussian,
     NaturalVariationalGaussian,
     VariationalGaussian,
     WhitenedVariationalGaussian,
@@ -118,6 +121,7 @@ def test_variational_gaussians(
     )
     likelihood = gpx.likelihoods.Gaussian(123)
     inducing_inputs = jnp.linspace(-5.0, 5.0, n_inducing).reshape(-1, 1)
+
     test_inputs = jnp.linspace(-5.0, 5.0, n_test).reshape(-1, 1)
 
     posterior = prior * likelihood
@@ -174,6 +178,61 @@ def test_variational_gaussians(
     assert sigma.shape == (n_test, n_test)
 
 
+@pytest.mark.parametrize("n_test", [10, 20])
+@pytest.mark.parametrize("n_inducing", [10, 20])
+@pytest.mark.parametrize(
+    "variational_family",
+    [
+        GraphVariationalGaussian,
+    ],
+)
+def test_graph_variational_gaussian(
+    n_test: int,
+    n_inducing: int,
+    variational_family: AbstractVariationalFamily,
+) -> None:
+    G = nx.barbell_graph(100, 0)
+    L = nx.laplacian_matrix(G).toarray()
+
+    kernel = gpx.kernels.GraphKernel(
+        laplacian=L,
+        lengthscale=2.3,
+        variance=3.2,
+        smoothness=6.1,
+    )
+    meanf = gpx.mean_functions.Constant()
+    prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
+    likelihood = gpx.likelihoods.Bernoulli(num_datapoints=G.number_of_nodes())
+
+    inducing_inputs = jnp.array(
+        np.random.randint(low=1, high=100, size=(n_inducing, 1))
+    ).astype(jnp.int64)
+
+    test_inputs = jnp.array(np.random.randint(low=0, high=1, size=(n_test, 1))).astype(
+        jnp.int64
+    )
+
+    posterior = prior * likelihood
+    q = variational_family(posterior=posterior, inducing_inputs=inducing_inputs)
+    # Test KL
+    kl = q.prior_kl()
+    assert isinstance(kl, jnp.ndarray)
+    assert kl.shape == ()
+    assert kl >= 0.0
+
+    # Test predictions
+    predictive_dist = q(test_inputs)
+    assert isinstance(predictive_dist, NumpyroDistribution)
+
+    mu = predictive_dist.mean
+    sigma = predictive_dist.covariance()
+
+    assert isinstance(mu, jnp.ndarray)
+    assert isinstance(sigma, jnp.ndarray)
+    assert mu.shape == (n_test,)
+    assert sigma.shape == (n_test, n_test)
+
+
 @pytest.mark.parametrize("n_test", [1, 10])
 @pytest.mark.parametrize("n_datapoints", [1, 10])
 @pytest.mark.parametrize("n_inducing", [1, 10, 20])