Merge pull request #95 from thomaspinder/Fix-conjugate-regression-bug-(single-datapoint)-

Fix conjugate regression bug for a single datapoint.

Author: thomaspinder

gpjax/gps.py

@@ -234,7 +234,8 @@ def mll(
 
             constant = jnp.array(-1.0) if negative else jnp.array(1.0)
             return constant * (
-                marginal_likelihood.log_prob(y.squeeze()).squeeze() + log_prior_density
+                marginal_likelihood.log_prob(jnp.atleast_1d(y.squeeze())).squeeze()
+                + log_prior_density
             )
 
         return mll

tests/test_abstractions.py

@@ -11,7 +11,7 @@
 tfd = tf.data
 
 
-@pytest.mark.parametrize("n", [20])
+@pytest.mark.parametrize("n", [1, 20])
 def test_fit(n):
     key = jr.PRNGKey(123)
     x = jnp.sort(jr.uniform(key=key, minval=-2.0, maxval=2.0, shape=(n, 1)), axis=0)
@@ -27,6 +27,7 @@ def test_fit(n):
     assert isinstance(optimised_params, dict)
     assert mll(optimised_params) < pre_mll_val
 
+
 def test_stop_grads():
     params = {"x": jnp.array(3.0), "y": jnp.array(4.0)}
     trainables = {"x": True, "y": False}

tests/test_gp.py

@@ -5,7 +5,7 @@
 import jax.random as jr
 import pytest
 
-from gpjax import Dataset, initialise, likelihoods, transform
+from gpjax import Dataset, initialise, transform
 from gpjax.gps import (
     AbstractGP,
     ConjugatePosterior,
@@ -36,7 +36,7 @@ def test_prior(num_datapoints):
     assert sigma.shape == (num_datapoints, num_datapoints)
 
 
-@pytest.mark.parametrize("num_datapoints", [2, 10])
+@pytest.mark.parametrize("num_datapoints", [1, 2, 10])
 def test_conjugate_posterior(num_datapoints):
     key = jr.PRNGKey(123)
     x = jnp.sort(
@@ -80,7 +80,7 @@ def test_conjugate_posterior(num_datapoints):
     assert sigma.shape == (num_datapoints, num_datapoints)
 
 
-@pytest.mark.parametrize("num_datapoints", [2, 10])
+@pytest.mark.parametrize("num_datapoints", [1, 2, 10])
 @pytest.mark.parametrize("likel", NonConjugateLikelihoods)
 def test_nonconjugate_posterior(num_datapoints, likel):
     key = jr.PRNGKey(123)