Support latest gpjax

Makefile

@@ -44,4 +44,9 @@ lint: sync
 test: sync ## Run the tests, start with the failing ones and break on first fail.
 	@$(UV_RUN) pytest -v -x --ff -rN -Wignore -s --tb=short --durations=0 --cov --cov-report=xml tests
 	@$(UV_RUN) pytest --nbmake --nbmake-kernel=python3 --durations=0 --nbmake-timeout=1000 --ignore=notebooks/frontends/GPJax.ipynb notebooks/
+	@if [ "$(UV_PYTHON)" = "python3.9" ]; then \
+		echo "Skipping GPJax notebook on python3.9"; \
+	else \
+		$(UV_RUN) pytest --nbmake --nbmake-kernel=python3 --durations=0 --nbmake-timeout=1000 notebooks/frontends/GPJax.ipynb; \
+	fi;
 	@echo -e "$(SUCCESS)Tests done$(RESET)"

docs/index.rst

@@ -200,16 +200,6 @@ To install JAX, follow `these instructions <https://github.com/google/jax#instal
 
    pip install gpjax
 
-.. warning::
-
-    .. raw:: html
-
-        <div style="color: var(--color-content-foreground);">
-
-    Currently, only some versions of `gpjax` are supported (we tested `gpjax==0.6.9`).
-
-    Furthermore, installation might be far from trivial and result in a broken environment. This is due to our conflicting dependencies, see https://github.com/JaxGaussianProcesses/GPJax/issues/441.
-
     .. raw:: html
 
         </div>

geometric_kernels/frontends/gpjax.py

@@ -10,10 +10,10 @@
 
 import gpjax
 import jax.numpy as jnp
-import tensorflow_probability.substrates.jax.bijectors as tfb
 from beartype.typing import List, TypeVar, Union
-from gpjax.base import param_field, static_field
+from flax import nnx
 from gpjax.kernels.computations.base import AbstractKernelComputation
+from gpjax.parameters import NonNegativeReal, PositiveReal
 from gpjax.typing import Array, ScalarFloat
 from jaxtyping import Float, Num
 
@@ -102,32 +102,55 @@ class GPJaxGeometricKernel(gpjax.kernels.AbstractKernel):
     :type variance: ScalarFloat
     """
 
-    nu: ScalarFloat = param_field(None, bijector=tfb.Softplus(), trainable=False)
-    lengthscale: Union[ScalarFloat, Float[Array, " D"]] = param_field(
-        None, bijector=tfb.Softplus()
-    )
-    variance: ScalarFloat = param_field(jnp.array(1.0), bijector=tfb.Softplus())
-    base_kernel: BaseGeometricKernel = static_field(None)
-    compute_engine: AbstractKernelComputation = static_field(
-        _GeometricKernelComputation(), repr=False
-    )
-    name: str = "Geometric Kernel"
+    nu: Union[ScalarFloat, nnx.Variable[ScalarFloat], None]
+    lengthscale: nnx.Variable[Union[ScalarFloat, Float[Array, " D"]]]
+    variance: nnx.Variable[ScalarFloat]
 
-    def __post_init__(self):
-        if self.base_kernel is None:
-            raise ValueError("base_kernel must be specified")
+    base_kernel: BaseGeometricKernel
+    compute_engine: AbstractKernelComputation = _GeometricKernelComputation()
+    name: str = "Geometric Kernel"
 
+    def __init__(
+        self,
+        base_kernel: BaseGeometricKernel,
+        lengthscale: Union[
+            Union[ScalarFloat, Float[Array, " D"]],
+            nnx.Variable[Union[ScalarFloat, Float[Array, " D"]]],
+            None,
+        ] = None,
+        nu: Union[ScalarFloat, nnx.Variable[ScalarFloat], None] = None,
+        variance: Union[ScalarFloat, nnx.Variable[ScalarFloat]] = 1.0,
+        trainable_nu: bool = False,
+    ):
+        active_dims = None
+        n_dims = None
+        super().__init__(active_dims, n_dims, self.compute_engine)
+
+        self.base_kernel = base_kernel
         default_params = self.base_kernel.init_params()
 
-        if self.nu is None:
-            self.nu = jnp.array(default_params["nu"])
-        if isinstance(self.nu, ScalarFloat):
-            self.nu = jnp.array([self.nu])
-
-        if self.lengthscale is None:
-            self.lengthscale = jnp.array(default_params["lengthscale"])
-        if isinstance(self.lengthscale, ScalarFloat):
-            self.lengthscale = jnp.array([self.lengthscale])
+        if lengthscale is None:
+            lengthscale = jnp.array(default_params["lengthscale"])
+        if nu is None:
+            nu = jnp.array(default_params["nu"])
+
+        if isinstance(lengthscale, nnx.Variable):
+            self.lengthscale = lengthscale
+        else:
+            self.lengthscale = PositiveReal(lengthscale)
+
+        self.trainable_nu = trainable_nu
+        if not trainable_nu:
+            self.nu = nu
+        elif isinstance(nu, nnx.Variable):
+            self.nu = nu
+        else:
+            self.nu = PositiveReal(nu)
+
+        if isinstance(variance, nnx.Variable):
+            self.variance = variance
+        else:
+            self.variance = NonNegativeReal(variance)
 
     @property
     def space(self) -> Union[Space, List[Space]]:
@@ -151,6 +174,7 @@ def __call__(
         :return:
             The N x M cross-covariance matrix.
         """
-        return self.variance * self.base_kernel.K(
-            {"lengthscale": self.lengthscale, "nu": self.nu}, x, y
+        nu_value = self.nu.value if self.trainable_nu else self.nu
+        return self.variance.value * self.base_kernel.K(
+            {"lengthscale": self.lengthscale.value, "nu": nu_value}, x, y
         )

pyproject.toml

@@ -68,7 +68,7 @@ default-groups = []
 dev = [
   # --- shared from test_requirements.txt ---
   "ipykernel",
-  "backends>=1.5.4",
+  "backends>=1.8.0",
   "plotly",
   "kaleido",
   "black==24.3.0",
@@ -94,18 +94,15 @@ dev = [
   "botorch>=0.9",
 
   # TensorFlow / GPflow / TFP split
-  'tensorflow==2.13.0; python_version=="3.8"',
-  'tensorflow==2.13.1; python_version=="3.9"',
-  'tensorflow==2.15;   python_version>="3.10" and python_version<"3.12"',
-  'tensorflow-probability==0.20.1; python_version < "3.10"',
-  'tensorflow-probability==0.23;   python_version >= "3.10" and python_version < "3.12"',
-  'gpflow==2.9.0; python_version < "3.10"',
-  'gpflow==2.9;    python_version >= "3.10" and python_version < "3.12"',
+  'tensorflow',
+  'tensorflow-probability',
+  'gpflow',
+  'tf_keras',
 
   # JAX family
   'jax',
   'jaxlib',
-  'jaxtyping==0.2.25; python_version == "3.9"',
-  'jaxtyping;         python_version != "3.9"',
+  'jaxtyping',
   'optax',
-]
+  'gpjax>=0.12.2; python_version >= "3.10" and python_version < "3.12"',  # gpjax is not supported on python-3.9 or older.
+]

tests/utils/test_utils.py

@@ -1,3 +1,5 @@
+import math
+
 import numpy as np
 import pytest
 
@@ -94,8 +96,8 @@ def test_hamming_distance(backend):
 def test_log_binomial(n):
     for k in range(n + 1):
         # Check that log_binomial gives the same result as the log of the
-        # binomial coefficient (as computed through `np.math.comb`).
-        assert np.isclose(np.log(np.math.comb(n, k)), log_binomial(n, k), atol=1e-10)
+        # binomial coefficient (as computed through `math.comb`).
+        assert np.isclose(np.log(math.comb(n, k)), log_binomial(n, k), atol=1e-10)
 
 
 @pytest.mark.parametrize("d", [0, 1, 2, 3, 5, 10])