Fix shape problem in GP priors

Shape asserts were added in existing tests that covered
both reparameterize=True/False.

Closes #5803

Author: michaelosthege

pymc/gp/gp.py

@@ -27,7 +27,6 @@
     JITTER_DEFAULT,
     cholesky,
     conditioned_vars,
-    infer_size,
     replace_with_values,
     solve_lower,
     solve_upper,
@@ -131,7 +130,7 @@ def _build_prior(self, name, X, reparameterize=True, jitter=JITTER_DEFAULT, **kw
         mu = self.mean_func(X)
         cov = stabilize(self.cov_func(X), jitter)
         if reparameterize:
-            size = infer_size(X, kwargs.pop("size", None))
+            size = np.shape(X)[0]
             v = pm.Normal(name + "_rotated_", mu=0.0, sigma=1.0, size=size, **kwargs)
             f = pm.Deterministic(name, mu + cholesky(cov).dot(v))
         else:
@@ -278,7 +277,7 @@ def _build_prior(self, name, X, reparameterize=True, jitter=JITTER_DEFAULT, **kw
         mu = self.mean_func(X)
         cov = stabilize(self.cov_func(X), jitter)
         if reparameterize:
-            size = infer_size(X, kwargs.pop("size", None))
+            size = np.shape(X)[0]
             v = pm.StudentT(name + "_rotated_", mu=0.0, sigma=1.0, nu=self.nu, size=size, **kwargs)
             f = pm.Deterministic(name, mu + cholesky(cov).dot(v))
         else:

pymc/gp/util.py

@@ -88,29 +88,6 @@ def replace_with_values(vars_needed, replacements=None, model=None):
     return fn(**replacements)
 
 
-def infer_size(X, n_points=None):
-    R"""
-    Maybe attempt to infer the size, or N, of a Gaussian process input matrix.
-
-    If a specific shape cannot be inferred, for instance if X is symbolic, then an
-    error is raised.
-
-    Parameters
-    ----------
-    X: array-like
-        Gaussian process input matrix.
-    n_points: None or int
-        The number of rows of `X`.  If `None`, the number of rows of `X` is
-        calculated from `X` if possible.
-    """
-    if n_points is None:
-        try:
-            n_points = int(X.shape[0])
-        except TypeError:
-            raise TypeError("Cannot infer 'shape', provide as an argument")
-    return n_points
-
-
 def stabilize(K, jitter=JITTER_DEFAULT):
     R"""
     Adds small diagonal to a covariance matrix.

pymc/tests/test_gp.py

@@ -824,6 +824,8 @@ def testLatent1(self):
             gp = pm.gp.Latent(mean_func=mean_func, cov_func=cov_func)
             f = gp.prior("f", self.X, reparameterize=False)
             p = gp.conditional("p", self.Xnew)
+        assert tuple(f.shape.eval()) == (self.X.shape[0],)
+        assert tuple(p.shape.eval()) == (self.Xnew.shape[0],)
         latent_logp = model.compile_logp()({"f": self.y, "p": self.pnew})
         npt.assert_allclose(latent_logp, self.logp, atol=0, rtol=1e-2)
 
@@ -834,6 +836,8 @@ def testLatent2(self):
             gp = pm.gp.Latent(mean_func=mean_func, cov_func=cov_func)
             f = gp.prior("f", self.X, reparameterize=True)
             p = gp.conditional("p", self.Xnew)
+        assert tuple(f.shape.eval()) == (self.X.shape[0],)
+        assert tuple(p.shape.eval()) == (self.Xnew.shape[0],)
         chol = np.linalg.cholesky(cov_func(self.X).eval())
         y_rotated = np.linalg.solve(chol, self.y - 0.5)
         latent_logp = model.compile_logp()({"f_rotated_": y_rotated, "p": self.pnew})
@@ -1068,6 +1072,8 @@ def testTPvsLatent(self):
             tp = pm.gp.TP(cov_func=cov_func, nu=self.nu)
             f = tp.prior("f", self.X, reparameterize=False)
             p = tp.conditional("p", self.Xnew)
+        assert tuple(f.shape.eval()) == (self.X.shape[0],)
+        assert tuple(p.shape.eval()) == (self.Xnew.shape[0],)
         tp_logp = model.compile_logp()({"f": self.y, "p": self.pnew})
         npt.assert_allclose(self.gp_latent_logp, tp_logp, atol=0, rtol=1e-2)
 
@@ -1077,6 +1083,8 @@ def testTPvsLatentReparameterized(self):
             tp = pm.gp.TP(cov_func=cov_func, nu=self.nu)
             f = tp.prior("f", self.X, reparameterize=True)
             p = tp.conditional("p", self.Xnew)
+        assert tuple(f.shape.eval()) == (self.X.shape[0],)
+        assert tuple(p.shape.eval()) == (self.Xnew.shape[0],)
         chol = np.linalg.cholesky(cov_func(self.X).eval())
         f_rotated = np.linalg.solve(chol, self.y)
         tp_logp = model.compile_logp()({"f_rotated_": f_rotated, "p": self.pnew})
@@ -1129,6 +1137,8 @@ def testLatentKronvsLatent(self):
             kron_gp = pm.gp.LatentKron(mean_func=self.mean, cov_funcs=self.cov_funcs)
             f = kron_gp.prior("f", self.Xs)
             p = kron_gp.conditional("p", self.Xnew)
+        assert tuple(f.shape.eval()) == (self.X.shape[0],)
+        assert tuple(p.shape.eval()) == (self.Xnew.shape[0],)
         kronlatent_logp = kron_model.compile_logp()({"f_rotated_": self.y_rotated, "p": self.pnew})
         npt.assert_allclose(kronlatent_logp, self.logp, atol=0, rtol=1e-3)
 
@@ -1186,6 +1196,8 @@ def testMarginalKronvsMarginalpredict(self):
             f = kron_gp.marginal_likelihood("f", self.Xs, self.y, sigma=self.sigma)
             p = kron_gp.conditional("p", self.Xnew)
             mu, cov = kron_gp.predict(self.Xnew)
+        assert tuple(f.shape.eval()) == (self.X.shape[0],)
+        assert tuple(p.shape.eval()) == (self.Xnew.shape[0],)
         npt.assert_allclose(mu, self.mu, atol=1e-5, rtol=1e-2)
         npt.assert_allclose(cov, self.cov, atol=1e-5, rtol=1e-2)
         with kron_model: