gpjax 0.9.1__tar.gz → 0.9.3__tar.gz

{gpjax-0.9.1 → gpjax-0.9.3}/PKG-INFO

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

{gpjax-0.9.1 → gpjax-0.9.3}/examples/intro_to_kernels.py

@@ -246,7 +246,7 @@ kernel = gpx.kernels.Matern52(
 prior = gpx.gps.Prior(mean_function=mean, kernel=kernel)
 
 likelihood = gpx.likelihoods.Gaussian(
-    num_datapoints=D.n, obs_stddev=PositiveReal(value=jnp.array(1e-3), tag="Static")
+    num_datapoints=D.n, obs_stdev=Static(jnp.array(1e-3))
 )  # Our function is noise-free, so we set the observation noise's standard deviation to a very small value
 
 no_opt_posterior = prior * likelihood

{gpjax-0.9.1 → gpjax-0.9.3}/examples/poisson.py

@@ -154,33 +154,22 @@ def logprob_fn(params):
     return gpx.objectives.log_posterior_density(model, D)
 
 
-# jit compile
-logprob_fn = jax.jit(logprob_fn)
-_ = logprob_fn(params)
+step_size = 1e-3
+inverse_mass_matrix = jnp.ones(53)
+nuts = blackjax.nuts(logprob_fn, step_size, inverse_mass_matrix)
 
+state = nuts.init(params)
 
-adapt = blackjax.window_adaptation(
-    blackjax.nuts, logprob_fn, num_adapt, target_acceptance_rate=0.65, progress_bar=True
-)
-
-# Initialise the chain
-last_state, kernel, _ = adapt.run(key, params)
-
-
-def inference_loop(rng_key, kernel, initial_state, num_samples):
-    def one_step(state, rng_key):
-        state, info = kernel(rng_key, state)
-        return state, (state, info)
-
-    keys = jax.random.split(rng_key, num_samples)
-    _, (states, infos) = jax.lax.scan(one_step, initial_state, keys, unroll=10)
+step = jax.jit(nuts.step)
 
-    return states, infos
 
+def one_step(state, rng_key):
+    state, info = step(rng_key, state)
+    return state, (state, info)
 
-# Sample from the posterior distribution
-states, infos = inference_loop(key, kernel, last_state, num_samples)
 
+keys = jax.random.split(key, num_samples)
+_, (states, infos) = jax.lax.scan(one_step, state, keys, unroll=10)
 
 # %% [markdown]
 # ### Sampler efficiency
@@ -190,7 +179,7 @@ states, infos = inference_loop(key, kernel, last_state, num_samples)
 # proposed sample, divided by the total number of steps run by the chain).
 
 # %%
-acceptance_rate = jnp.mean(infos.acceptance_probability)
+acceptance_rate = jnp.mean(infos.acceptance_rate)
 print(f"Acceptance rate: {acceptance_rate:.2f}")
 
 # %%

{gpjax-0.9.1 → gpjax-0.9.3}/examples/uncollapsed_vi.py

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.4
+#       jupytext_version: 1.11.2
 #   kernelspec:
 #     display_name: gpjax_beartype
 #     language: python
@@ -319,7 +319,6 @@ opt_rep, history = gpx.fit(
     model=q,
     objective=lambda p, d: -gpx.objectives.elbo(p, d),
     train_data=D,
-    params_bijection=params_bijection,
     optim=ox.adam(learning_rate=0.01),
     num_iters=3000,
     key=jr.key(42),

{gpjax-0.9.1 → gpjax-0.9.3}/gpjax/__init__.py

@@ -40,7 +40,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.9.1"
+__version__ = "0.9.3"
 
 __all__ = [
     "base",

{gpjax-0.9.1 → gpjax-0.9.3}/gpjax/gps.py

@@ -17,6 +17,7 @@ from abc import abstractmethod
 
 import beartype.typing as tp
 from cola.annotations import PSD
+from cola.linalg.algorithm_base import Algorithm
 from cola.linalg.decompositions.decompositions import Cholesky
 from cola.linalg.inverse.inv import solve
 from cola.ops.operators import I_like
@@ -530,6 +531,7 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
         train_data: Dataset,
         key: KeyArray,
         num_features: int | None = 100,
+        solver_algorithm: tp.Optional[Algorithm] = Cholesky(),
     ) -> FunctionalSample:
         r"""Draw approximate samples from the Gaussian process posterior.
 
@@ -563,6 +565,11 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
             key (KeyArray): The random seed used for the sample(s).
             num_features (int): The number of features used when approximating the
                 kernel.
+            solver_algorithm (Optional[Algorithm], optional): The algorithm to use for the solves of
+                the inverse of the covariance matrix. See the
+                [CoLA documentation](https://cola.readthedocs.io/en/latest/package/cola.linalg.html#algorithms)
+                for which solver to pick. For PSD matrices, CoLA currently recommends Cholesky() for small
+                matrices and CG() for larger matrices. Select Auto() to let CoLA decide. Defaults to Cholesky().
 
         Returns:
             FunctionalSample: A function representing an approximate sample from the Gaussian
@@ -588,7 +595,7 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
         canonical_weights = solve(
             Sigma,
             y + eps - jnp.inner(Phi, fourier_weights),
-            Cholesky(),
+            solver_algorithm,
         )  #  [N, B]
 
         def sample_fn(test_inputs: Float[Array, "n D"]) -> Float[Array, "n B"]:

{gpjax-0.9.1 → gpjax-0.9.3}/gpjax/likelihoods.py

@@ -28,7 +28,6 @@ from gpjax.integrators import (
     GHQuadratureIntegrator,
 )
 from gpjax.parameters import (
-    Parameter,
     PositiveReal,
     Static,
 )
@@ -152,10 +151,9 @@ class Gaussian(AbstractLikelihood):
                 likelihoods. Must be an instance of `AbstractIntegrator`. For the Gaussian likelihood, this defaults to
                 the `AnalyticalGaussianIntegrator`, as the expected log likelihood can be computed analytically.
         """
-        if isinstance(obs_stddev, Parameter):
-            self.obs_stddev = obs_stddev
-        else:
-            self.obs_stddev = PositiveReal(jnp.asarray(obs_stddev))
+        if not isinstance(obs_stddev, (PositiveReal, Static)):
+            obs_stddev = PositiveReal(jnp.asarray(obs_stddev))
+        self.obs_stddev = obs_stddev
 
         super().__init__(num_datapoints, integrator)
 

{gpjax-0.9.1 → gpjax-0.9.3}/gpjax/scan.py

@@ -22,7 +22,6 @@ from beartype.typing import (
 )
 import jax
 from jax import lax
-from jax.experimental import host_callback as hcb
 import jax.numpy as jnp
 import jax.tree_util as jtu
 from jaxtyping import (
@@ -54,7 +53,8 @@ def _callback(cond: ScalarBool, func: Callable, *args: Any) -> None:
 
     def _do_callback(_) -> int:
         """Perform the callback."""
-        return hcb.id_tap(func, *args, result=_dummy_result)
+        jax.debug.callback(func, *args)
+        return _dummy_result
 
     def _not_callback(_) -> int:
         """Do nothing."""
@@ -113,19 +113,19 @@ def vscan(
     _progress_bar = trange(_length)
     _progress_bar.set_description("Compiling...", refresh=True)
 
-    def _set_running(args: Any, transform: Any) -> None:
+    def _set_running(*args: Any) -> None:
         """Set the tqdm progress bar to running."""
         _progress_bar.set_description("Running", refresh=False)
 
-    def _update_tqdm(args: Any, transform: Any) -> None:
+    def _update_tqdm(*args: Any) -> None:
         """Update the tqdm progress bar with the latest objective value."""
         _value, _iter_num = args
-        _progress_bar.update(_iter_num)
+        _progress_bar.update(_iter_num.item())
 
         if log_value and _value is not None:
             _progress_bar.set_postfix({"Value": f"{_value: .2f}"})
 
-    def _close_tqdm(args: Any, transform: Any) -> None:
+    def _close_tqdm(*args: Any) -> None:
         """Close the tqdm progress bar."""
         _progress_bar.close()
 
@@ -145,16 +145,16 @@ def vscan(
         _is_last: bool = iter_num == _length - 1
 
         # Update progress bar, if first of log_rate.
-        _callback(_is_first, _set_running, (y, log_rate))
+        _callback(_is_first, _set_running)
 
         # Update progress bar, if multiple of log_rate.
-        _callback(_is_multiple, _update_tqdm, (y, log_rate))
+        _callback(_is_multiple, _update_tqdm, y, log_rate)
 
         # Update progress bar, if remainder.
-        _callback(_is_remainder, _update_tqdm, (y, _remainder))
+        _callback(_is_remainder, _update_tqdm, y, _remainder)
 
         # Close progress bar, if last iteration.
-        _callback(_is_last, _close_tqdm, (y, None))
+        _callback(_is_last, _close_tqdm)
 
         return carry, y
 

{gpjax-0.9.1 → gpjax-0.9.3}/gpjax/variational_families.py

@@ -108,10 +108,17 @@ class AbstractVariationalGaussian(AbstractVariationalFamily[L]):
     def __init__(
         self,
         posterior: AbstractPosterior[P, L],
-        inducing_inputs: Float[Array, "N D"],
+        inducing_inputs: tp.Union[
+            Float[Array, "N D"],
+            Real,
+            Static,
+        ],
         jitter: ScalarFloat = 1e-6,
     ):
-        self.inducing_inputs = Static(inducing_inputs)
+        if not isinstance(inducing_inputs, (Real, Static)):
+            inducing_inputs = Real(inducing_inputs)
+
+        self.inducing_inputs = inducing_inputs
         self.jitter = jitter
 
         super().__init__(posterior)

{gpjax-0.9.1 → gpjax-0.9.3}/mkdocs.yml

@@ -24,8 +24,8 @@ nav:
     - Barycentres: _examples/barycentres.md
     - Deep kernel learning: _examples/deep_kernels.md
     - Graph kernels: _examples/graph_kernels.md
-    - Sparse GPs: _examples/uncollapsed_vi.md
-    - Stochastic sparse GPs: _examples/collapsed_vi.md
+    - Sparse GPs: _examples/collapsed_vi.md
+    - Stochastic sparse GPs: _examples/uncollapsed_vi.md
     - Bayesian Optimisation: _examples/bayesian_optimisation.md
     - Decision Making: _examples/decision_making.md
     - Multi-output GPs for Ocean Modelling: _examples/oceanmodelling.md

{gpjax-0.9.1 → gpjax-0.9.3}/tests/test_gps.py

@@ -25,13 +25,15 @@ from typing import (
     Type,
 )
 
+from cola.linalg.algorithm_base import Auto
+from cola.linalg.decompositions.decompositions import Cholesky
+from cola.linalg.inverse.cg import CG
 from jax import config
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
 import tensorflow_probability.substrates.jax.distributions as tfd
 
-# from gpjax.dataset import Dataset
 from gpjax.dataset import Dataset
 from gpjax.distributions import GaussianDistribution
 from gpjax.gps import (
@@ -283,7 +285,10 @@ def test_prior_sample_approx(num_datapoints, kernel, mean_function):
 @pytest.mark.parametrize("num_datapoints", [1, 5])
 @pytest.mark.parametrize("kernel", [RBF, Matern52])
 @pytest.mark.parametrize("mean_function", [Zero, Constant])
-def test_conjugate_posterior_sample_approx(num_datapoints, kernel, mean_function):
+@pytest.mark.parametrize("solver_algorithm", [Cholesky(), CG(), Auto()])
+def test_conjugate_posterior_sample_approx(
+    num_datapoints, kernel, mean_function, solver_algorithm
+):
     kern = kernel(lengthscale=jnp.array([5.0, 1.0]), variance=0.1)
     p = Prior(kernel=kern, mean_function=mean_function()) * Gaussian(
         num_datapoints=num_datapoints
@@ -310,26 +315,28 @@ def test_conjugate_posterior_sample_approx(num_datapoints, kernel, mean_function
     # with pytest.raises(ValidationErrors):
     #     p.sample_approx(1, D, key, 0.5)
 
-    sampled_fn = p.sample_approx(1, D, key, 100)
+    sampled_fn = p.sample_approx(1, D, key, 100, solver_algorithm=solver_algorithm)
     assert isinstance(sampled_fn, Callable)  # check type
 
     x = jr.uniform(key=key, minval=-2.0, maxval=2.0, shape=(num_datapoints, 2))
     evals = sampled_fn(x)
     assert evals.shape == (num_datapoints, 1.0)  # check shape
 
-    sampled_fn_2 = p.sample_approx(1, D, key, 100)
+    sampled_fn_2 = p.sample_approx(1, D, key, 100, solver_algorithm=solver_algorithm)
     evals_2 = sampled_fn_2(x)
     max_delta = jnp.max(jnp.abs(evals - evals_2))
     assert max_delta == 0.0  # samples same for same seed
 
     new_key = jr.key(12345)
-    sampled_fn_3 = p.sample_approx(1, D, new_key, 100)
+    sampled_fn_3 = p.sample_approx(
+        1, D, new_key, 100, solver_algorithm=solver_algorithm
+    )
     evals_3 = sampled_fn_3(x)
     max_delta = jnp.max(jnp.abs(evals - evals_3))
     assert max_delta > 0.01  # samples different for different seed
 
     # Check validty of samples using Monte-Carlo
-    sampled_fn = p.sample_approx(10_000, D, key, 100)
+    sampled_fn = p.sample_approx(10_000, D, key, 100, solver_algorithm=solver_algorithm)
     sampled_evals = sampled_fn(x)
     approx_mean = jnp.mean(sampled_evals, -1)
     approx_var = jnp.var(sampled_evals, -1)

gpjax-0.9.1/.github/workflows/labeler.yml

@@ -1,18 +0,0 @@
-name: Labeler
-
-on:
-  push:
-    branches:
-      - main
-
-jobs:
-  labeler:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check out the repository
-        uses: actions/checkout@v3.5.2
-
-      - name: Run Labeler
-        uses: crazy-max/ghaction-github-labeler@v4.1.0
-        with:
-          skip-delete: true