Merge pull request #179 from JaxGaussianProcesses/variational_expecta…

…ion_quadrature_bug Update variational_inference.py
JaxGaussianProcesses · Jan 9, 2023 · 3a40bac · 3a40bac
2 parents c0809b4 + 42cae0a
commit 3a40bac
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gpjax/variational_inference.py b/gpjax/variational_inference.py
@@ -153,7 +153,7 @@ def q_moments(x):
         log_prob = vmap(lambda f, y: link_function(params["likelihood"], f).log_prob(y))
 
         # ≈ ∫[log(p(y|f(x))) q(f(x))] df(x)
-        expectation = gauss_hermite_quadrature(log_prob, mean, variance, y=y)
+        expectation = gauss_hermite_quadrature(log_prob, mean, jnp.sqrt(variance), y=y)
 
         return expectation