fixing dtype issues

dfm · Mar 26, 2021 · bdea294 · bdea294
1 parent fcc0919
commit bdea294
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/src/tinygp/gp.py b/src/tinygp/gp.py
@@ -11,8 +11,6 @@
 from .kernels import Kernel
 from .means import Mean, constant_mean, zero_mean
 
-LOG2PI = jnp.log(2 * jnp.pi)
-
 
 class GaussianProcess:
     def __init__(
@@ -49,10 +47,8 @@ def __init__(
         )
         self.lower = lower
         self.chol = linalg.cholesky(self.K, lower=self.lower)
-        self.norm = (
-            jnp.sum(jnp.log(jnp.diag(self.chol)))
-            + 0.5 * self.X.shape[0] * LOG2PI
-        )
+        self.norm = jnp.sum(jnp.log(jnp.diag(self.chol)))
+        self.norm += 0.5 * self.X.shape[0] * jnp.log(2 * jnp.pi)
 
     def condition(self, y: jnp.ndarray) -> jnp.ndarray:
         self.y = jnp.broadcast_to(y, (self.size,))