Fix AnalyticalGuassianIntegrator

gpjax/integrators.py

@@ -125,7 +125,8 @@ class AnalyticalGaussianIntegrator(AbstractIntegrator):
     variational distribution $`q(f) = \mathcal{N}(f|m, s)`$, the expected
     log-likelihood is given by
     ```math
-    \mathbb{E}_{q(f)}[\log p(y|f)] = -\frac{1}{2}\left(\log(2\pi\sigma^2) + \frac{1}{\sigma^2}\mathbb{E}_{q(f)}[(y-f)^2 + s]\right)
+    \mathbb{E}_{q(f)}[\log p(y|f)] = -\frac{1}{2}\left(\log(2\pi\sigma^2) + \frac{1}{\sigma^2}((y-m)^2 + s)\right)
+    ```
     """
 
     def integrate(
@@ -153,7 +154,8 @@ def integrate(
         sq_error = jnp.square(y - mean)
         log2pi = jnp.log(2.0 * jnp.pi)
         val = jnp.sum(
-            log2pi + jnp.log(obs_stddev) + (sq_error + variance) / obs_stddev, axis=1
+            log2pi + jnp.log(obs_stddev**2) + (sq_error + variance) / obs_stddev**2,
+            axis=1
         )
         return -0.5 * val
 

tests/test_integrators.py

@@ -57,7 +57,9 @@ def fun(x, y):
 
 
 @pytest.mark.parametrize("jit", [True, False])
-@pytest.mark.parametrize("params", [(0.5, -2.57236494), (1.0, -1.91893853)])
+@pytest.mark.parametrize(
+    "params", [(0.5, -4.22579135), (1.0, -1.91893853), (0.01, -9996.31376835)]
+)
 def test_analytical_gaussian(jit: bool, params: tp.Tuple[float, float]):
     obs_stddev, expected = params
     likelihood = Gaussian(num_datapoints=1, obs_stddev=jnp.array([obs_stddev]))