gpjax 0.10.1__tar.gz → 0.10.2__tar.gz

{gpjax-0.10.1 → gpjax-0.10.2}/PKG-INFO

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

{gpjax-0.10.1 → gpjax-0.10.2}/docs/sharp_bits.md

@@ -175,3 +175,60 @@ mini-batch optimisation of the parameters of your sparse Gaussian process model.
 model will scale linearly in the batch size and quadratically in the number of inducing
 points. We demonstrate its use in
 [our sparse stochastic variational inference notebook](_examples/uncollapsed_vi.md).
+
+## JIT compilation
+
+There are a subset of operations in GPJax that are not JIT compatible by default. This
+is because we have assertions in place to check the properties of the parameters. For
+example, we check that the lengthscale parameter that a user provides is positive. This
+makes for a better user experience as we can provide more informative error messages;
+however, JIT compiling functions wherein these assertions are made will break the code.
+As an example, consider the following code:
+
+```python
+import jax
+import jax.numpy as jnp
+import gpjax as gpx
+
+x = jnp.linspace(0, 1, 10)[:, None]
+
+def compute_gram(lengthscale):
+    k = gpx.kernels.RBF(active_dims=[0], lengthscale=lengthscale, variance=jnp.array(1.0))
+    return k.gram(x)
+
+compute_gram(1.0)
+```
+
+so far so good. However, if we try to JIT compile this function, we will get an error:
+
+```python
+jit_compute_gram = jax.jit(compute_gram)
+try:
+    jit_compute_gram(1.0)
+except Exception as e:
+    print(e)
+```
+
+This error is due to the fact that the `RBF` kernel contains an assertion that checks
+that the lengthscale is positive. It does not matter that the assertion is satisfied;
+the very presence of the assertion will break JIT compilation.
+
+To resolve this, we can use the `checkify` decorator to remove the assertion. This will
+allow the function to be JIT compiled.
+
+```python
+from jax.experimental import checkify
+
+jit_compute_gram = jax.jit(checkify.checkify(compute_gram))
+error, value = jit_compute_gram(1.0)
+```
+By virtue of the `checkify.checkify`, a tuple is returned where the first element is the
+output of the assertion, and the second element is the value of the function. 
+
+This design is not perfect, and in an ideal world we would not enforce the user to wrap
+their code in `checkify.checkify`. We are actively looking into cleaner ways to provide
+guardrails in a less intrusive manner. However, for now, should you try to JIT compile
+a component of GPJax wherein there is an assertion, you will need to wrap the function
+in `checkify.checkify` as shown above.
+
+For more on `checkify`, please see the [JAX Checkify Doc](https://docs.jax.dev/en/latest/debugging/checkify_guide.html).

{gpjax-0.10.1 → gpjax-0.10.2}/examples/oak_example.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.11.2
+#       jupytext_version: 1.16.7
 #   kernelspec:
 #     display_name: docs
 #     language: python
@@ -37,19 +37,21 @@
 # %%
 import jax
 from jax import config
-
-config.update("jax_enable_x64", True)  # Enable Float64 precision
-
 import jax.numpy as jnp
-import matplotlib.pyplot as plt
 from matplotlib.colors import ListedColormap
+import matplotlib.pyplot as plt
 import optax
 
 import gpjax as gpx
 from gpjax.dataset import Dataset
-from gpjax.kernels import OrthogonalAdditiveKernel, RBF
+from gpjax.kernels import (
+    RBF,
+    OrthogonalAdditiveKernel,
+)
 from gpjax.typing import KeyArray
 
+config.update("jax_enable_x64", True)  # Enable Float64 precision
+
 
 # %%
 def f(x: jnp.ndarray) -> jnp.ndarray:
@@ -198,7 +200,7 @@ def main():
 
     print("\nRelative Importance of Input Dimensions:")
     for i, imp in enumerate(relative_importance):
-        print(f"Dimension {i+1}: {imp:.4f}")
+        print(f"Dimension {i + 1}: {imp:.4f}")
 
     if opt_posterior.params.kernel.coeffs_2 is not None:
         # Analyze second-order interactions

{gpjax-0.10.1 → gpjax-0.10.2}/gpjax/__init__.py

@@ -39,7 +39,7 @@ __license__ = "MIT"
 __description__ = "Didactic Gaussian processes in JAX"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.10.1"
+__version__ = "0.10.2"
 
 __all__ = [
     "base",

{gpjax-0.10.1 → gpjax-0.10.2}/gpjax/kernels/nonstationary/polynomial.py

@@ -46,7 +46,7 @@ class Polynomial(AbstractKernel):
         self,
         active_dims: tp.Union[list[int], slice, None] = None,
         degree: int = 2,
-        shift: tp.Union[ScalarFloat, nnx.Variable[ScalarArray]] = 0.0,
+        shift: tp.Union[ScalarFloat, nnx.Variable[ScalarArray]] = 1.0,
         variance: tp.Union[ScalarFloat, nnx.Variable[ScalarArray]] = 1.0,
         n_dims: tp.Union[int, None] = None,
         compute_engine: AbstractKernelComputation = DenseKernelComputation(),

{gpjax-0.10.1 → gpjax-0.10.2}/gpjax/parameters.py

@@ -1,6 +1,7 @@
 import typing as tp
 
 from flax import nnx
+from jax.experimental import checkify
 import jax.numpy as jnp
 import jax.tree_util as jtu
 from jax.typing import ArrayLike
@@ -84,8 +85,7 @@ class PositiveReal(Parameter[T]):
 
     def __init__(self, value: T, tag: ParameterTag = "positive", **kwargs):
         super().__init__(value=value, tag=tag, **kwargs)
-
-        _check_is_positive(self.value)
+        _safe_assert(_check_is_positive, self.value)
 
 
 class Real(Parameter[T]):
@@ -101,7 +101,17 @@ class SigmoidBounded(Parameter[T]):
     def __init__(self, value: T, tag: ParameterTag = "sigmoid", **kwargs):
         super().__init__(value=value, tag=tag, **kwargs)
 
-        _check_in_bounds(self.value, 0.0, 1.0)
+        # Only perform validation in non-JIT contexts
+        if (
+            not isinstance(value, jnp.ndarray)
+            or not getattr(value, "aval", None) is None
+        ):
+            _safe_assert(
+                _check_in_bounds,
+                self.value,
+                low=jnp.array(0.0),
+                high=jnp.array(1.0),
+            )
 
 
 class Static(nnx.Variable[T]):
@@ -120,8 +130,13 @@ class LowerTriangular(Parameter[T]):
     def __init__(self, value: T, tag: ParameterTag = "lower_triangular", **kwargs):
         super().__init__(value=value, tag=tag, **kwargs)
 
-        _check_is_square(self.value)
-        _check_is_lower_triangular(self.value)
+        # Only perform validation in non-JIT contexts
+        if (
+            not isinstance(value, jnp.ndarray)
+            or not getattr(value, "aval", None) is None
+        ):
+            _safe_assert(_check_is_square, self.value)
+            _safe_assert(_check_is_lower_triangular, self.value)
 
 
 DEFAULT_BIJECTION = {
@@ -132,36 +147,83 @@ DEFAULT_BIJECTION = {
 }
 
 
-def _check_is_arraylike(value: T):
+def _check_is_arraylike(value: T) -> None:
+    """Check if a value is array-like.
+
+    Args:
+        value: The value to check.
+
+    Raises:
+        TypeError: If the value is not array-like.
+    """
     if not isinstance(value, (ArrayLike, list)):
         raise TypeError(
             f"Expected parameter value to be an array-like type. Got {value}."
         )
 
 
-def _check_is_positive(value: T):
-    if jnp.any(value < 0):
-        raise ValueError(
-            f"Expected parameter value to be strictly positive. Got {value}."
-        )
+@checkify.checkify
+def _check_is_positive(value):
+    checkify.check(
+        jnp.all(value > 0), "value needs to be positive, got {value}", value=value
+    )
 
 
-def _check_is_square(value: T):
-    if value.shape[0] != value.shape[1]:
-        raise ValueError(
-            f"Expected parameter value to be a square matrix. Got {value}."
-        )
+@checkify.checkify
+def _check_is_square(value: T) -> None:
+    """Check if a value is a square matrix.
 
+    Args:
+        value: The value to check.
 
-def _check_is_lower_triangular(value: T):
-    if not jnp.all(jnp.tril(value) == value):
-        raise ValueError(
-            f"Expected parameter value to be a lower triangular matrix. Got {value}."
-        )
+    Raises:
+        ValueError: If the value is not a square matrix.
+    """
+    checkify.check(
+        value.shape[0] == value.shape[1],
+        "value needs to be a square matrix, got {value}",
+        value=value,
+    )
 
 
-def _check_in_bounds(value: T, low: float, high: float):
-    if jnp.any((value < low) | (value > high)):
-        raise ValueError(
-            f"Expected parameter value to be bounded between {low} and {high}. Got {value}."
-        )
+@checkify.checkify
+def _check_is_lower_triangular(value: T) -> None:
+    """Check if a value is a lower triangular matrix.
+
+    Args:
+        value: The value to check.
+
+    Raises:
+        ValueError: If the value is not a lower triangular matrix.
+    """
+    checkify.check(
+        jnp.all(jnp.tril(value) == value),
+        "value needs to be a lower triangular matrix, got {value}",
+        value=value,
+    )
+
+
+@checkify.checkify
+def _check_in_bounds(value: T, low: T, high: T) -> None:
+    """Check if a value is bounded between low and high.
+
+    Args:
+        value: The value to check.
+        low: The lower bound.
+        high: The upper bound.
+
+    Raises:
+        ValueError: If any element of value is outside the bounds.
+    """
+    checkify.check(
+        jnp.all((value >= low) & (value <= high)),
+        "value needs to be bounded between {low} and {high}, got {value}",
+        value=value,
+        low=low,
+        high=high,
+    )
+
+
+def _safe_assert(fn: tp.Callable[[tp.Any], None], value: T, **kwargs) -> None:
+    error, _ = fn(value, **kwargs)
+    checkify.check_error(error)

gpjax-0.10.2/tests/test_parameters.py

@@ -0,0 +1,113 @@
+from flax import nnx
+from jax import jit
+from jax.experimental import checkify
+import jax.numpy as jnp
+import pytest
+
+from gpjax.parameters import (
+    DEFAULT_BIJECTION,
+    LowerTriangular,
+    Parameter,
+    PositiveReal,
+    Real,
+    SigmoidBounded,
+    Static,
+    _check_in_bounds,
+    _check_is_lower_triangular,
+    _check_is_positive,
+    _check_is_square,
+    _safe_assert,
+    transform,
+)
+
+
+@pytest.mark.parametrize(
+    "param, value",
+    [
+        (PositiveReal, 1.0),
+        (Real, 2.0),
+        (SigmoidBounded, 0.5),
+    ],
+)
+def test_transform(param, value):
+    # Create mock parameters and bijectors
+    params = nnx.State(
+        {
+            "param1": param(value),
+            "param2": Parameter(2.0, tag="real"),
+        }
+    )
+
+    # Test forward transformation
+    t_params = transform(params, DEFAULT_BIJECTION)
+    t_param1_expected = DEFAULT_BIJECTION[params["param1"]._tag].forward(value)
+    assert jnp.allclose(t_params["param1"].value, t_param1_expected)
+    assert jnp.allclose(t_params["param2"].value, 2.0)
+
+    # Test inverse transformation
+    it_params = transform(t_params, DEFAULT_BIJECTION, inverse=True)
+    assert repr(it_params) == repr(params)
+
+
+@pytest.mark.parametrize(
+    "param, tag",
+    [
+        (PositiveReal(1.0), "positive"),
+        (Real(2.0), "real"),
+        (SigmoidBounded(0.5), "sigmoid"),
+        (Static(2.0), "static"),
+        (LowerTriangular(jnp.eye(2)), "lower_triangular"),
+    ],
+)
+def test_default_tags(param, tag):
+    assert param._tag == tag
+
+
+def test_check_is_positive():
+    # Check singleton
+    _safe_assert(_check_is_positive, jnp.array(3.0))
+    # Check array
+    _safe_assert(_check_is_positive, jnp.array([3.0, 4.0]))
+
+    # Check negative singleton
+    with pytest.raises(ValueError):
+        _safe_assert(_check_is_positive, jnp.array(-3.0))
+
+    # Check negative array
+    with pytest.raises(ValueError):
+        _safe_assert(_check_is_positive, jnp.array([-3.0, 4.0]))
+
+    # Test that functions wrapping _check_is_positive are jittable
+    def _dummy_fn(value):
+        _safe_assert(_check_is_positive, value)
+
+    jitted_fn = jit(checkify.checkify(_dummy_fn))
+    jitted_fn(jnp.array(3.0))
+
+
+def test_check_is_square():
+    # Check square matrix
+    _safe_assert(_check_is_square, jnp.full((2, 2), 1.0))
+    # Check non-square matrix
+    with pytest.raises(ValueError):
+        _safe_assert(_check_is_square, jnp.full((2, 3), 1.0))
+
+
+def test_check_is_lower_triangular():
+    # Check lower triangular matrix
+    _safe_assert(_check_is_lower_triangular, jnp.tril(jnp.eye(2)))
+    # Check non-lower triangular matrix
+    with pytest.raises(ValueError):
+        _safe_assert(_check_is_lower_triangular, jnp.linspace(0.0, 1.0, 4))
+
+
+def test_check_in_bounds():
+    # Check in bounds
+    _safe_assert(
+        _check_in_bounds, jnp.array(0.5), low=jnp.array(0.0), high=jnp.array(1.0)
+    )
+    # Check out of bounds
+    with pytest.raises(ValueError):
+        _safe_assert(
+            _check_in_bounds, jnp.array(1.5), low=jnp.array(0.0), high=jnp.array(1.0)
+        )